From 3de904c7ea1852d8cf4c1c75f94e803e7c54f3e1 Mon Sep 17 00:00:00 2001 From: Buqian Zheng Date: Mon, 28 Apr 2025 10:52:40 +0800 Subject: [PATCH] feat: add cachinglayer to sealed segment (#41436) issue: https://github.com/milvus-io/milvus/issues/41435 --------- Signed-off-by: Buqian Zheng --- .gitignore | 4 +- configs/milvus.yaml | 6 +- docs/design_docs/segcore/segment_sealed.md | 1 - internal/core/conanfile.py | 2 +- internal/core/src/CMakeLists.txt | 2 + internal/core/src/cachinglayer/CMakeLists.txt | 13 + internal/core/src/cachinglayer/CacheSlot.h | 372 ++++++ internal/core/src/cachinglayer/Manager.cpp | 62 + internal/core/src/cachinglayer/Manager.h | 66 ++ internal/core/src/cachinglayer/Translator.h | 60 + internal/core/src/cachinglayer/Utils.h | 300 +++++ .../core/src/cachinglayer/lrucache/DList.cpp | 201 ++++ .../core/src/cachinglayer/lrucache/DList.h | 108 ++ .../src/cachinglayer/lrucache/ListNode.cpp | 211 ++++ .../core/src/cachinglayer/lrucache/ListNode.h | 183 +++ internal/core/src/common/ArrowDataWrapper.h | 39 + internal/core/src/common/Channel.h | 3 +- internal/core/src/common/Chunk.h | 6 + internal/core/src/common/ChunkWriter.cpp | 20 +- internal/core/src/common/ChunkWriter.h | 20 +- internal/core/src/common/FieldData.h | 22 +- internal/core/src/common/IndexMeta.h | 4 + internal/core/src/common/LoadInfo.h | 2 - internal/core/src/common/Span.h | 6 +- internal/core/src/common/Types.h | 2 +- .../src/exec/expression/BinaryRangeExpr.cpp | 3 +- .../core/src/exec/expression/CompareExpr.h | 28 +- internal/core/src/exec/expression/Expr.cpp | 1 + internal/core/src/exec/expression/Expr.h | 59 +- .../src/exec/expression/JsonContainsExpr.cpp | 18 +- .../core/src/exec/expression/TermExpr.cpp | 4 +- .../core/src/exec/expression/UnaryExpr.cpp | 156 ++- .../core/src/exec/operator/FilterBitsNode.cpp | 2 + .../src/exec/operator/IterativeFilterNode.cpp | 2 + .../src/exec/operator/RandomSampleNode.cpp | 3 +- .../src/exec/operator/VectorSearchNode.cpp | 2 + .../operator/groupby/SearchGroupByOperator.h | 34 +- internal/core/src/futures/Executor.cpp | 2 +- .../src/index/JsonKeyStatsInvertedIndex.h | 16 +- internal/core/src/index/SkipIndex.h | 13 +- internal/core/src/mmap/ChunkData.h | 4 +- internal/core/src/mmap/ChunkVector.h | 2 + internal/core/src/mmap/ChunkedColumn.h | 402 +++---- internal/core/src/mmap/Column.h | 137 --- internal/core/src/mmap/Types.h | 15 +- .../core/src/monitor/prometheus_client.cpp | 214 ++++ internal/core/src/monitor/prometheus_client.h | 68 ++ .../core/src/query/CachedSearchIterator.cpp | 25 +- .../core/src/query/CachedSearchIterator.h | 8 +- internal/core/src/query/SearchOnSealed.cpp | 40 +- internal/core/src/query/SearchOnSealed.h | 36 +- .../src/segcore/ChunkedSegmentSealedImpl.cpp | 1039 +++++------------ .../src/segcore/ChunkedSegmentSealedImpl.h | 128 +- internal/core/src/segcore/ConcurrentVector.h | 52 - internal/core/src/segcore/DeletedRecord.h | 31 +- internal/core/src/segcore/FieldIndexing.h | 8 +- internal/core/src/segcore/InsertRecord.h | 368 +++--- internal/core/src/segcore/Record.h | 2 + .../core/src/segcore/SegmentChunkReader.cpp | 77 +- .../core/src/segcore/SegmentGrowingImpl.cpp | 24 +- .../core/src/segcore/SegmentGrowingImpl.h | 37 +- .../core/src/segcore/SegmentInterface.cpp | 1 + internal/core/src/segcore/SegmentInterface.h | 79 +- internal/core/src/segcore/SegmentSealed.h | 22 +- internal/core/src/segcore/Utils.h | 10 +- .../core/src/segcore/load_field_data_c.cpp | 6 - internal/core/src/segcore/load_field_data_c.h | 3 - internal/core/src/segcore/reduce/Reduce.cpp | 3 +- internal/core/src/segcore/segcore_init_c.cpp | 9 + internal/core/src/segcore/segcore_init_c.h | 5 + internal/core/src/segcore/segment_c.cpp | 60 +- internal/core/src/segcore/segment_c.h | 9 +- .../storagev1translator/ChunkTranslator.cpp | 166 +++ .../storagev1translator/ChunkTranslator.h | 71 ++ .../DefaultValueChunkTranslator.cpp | 96 ++ .../DefaultValueChunkTranslator.h | 62 + .../InsertRecordTranslator.cpp | 139 +++ .../InsertRecordTranslator.h | 72 ++ internal/core/src/storage/ChunkCache.cpp | 202 ---- internal/core/src/storage/ChunkCache.h | 83 -- internal/core/src/storage/Event.cpp | 4 +- internal/core/src/storage/MmapManager.h | 24 +- internal/core/src/storage/ThreadPool.h | 1 - internal/core/src/storage/Util.cpp | 62 +- internal/core/src/storage/Util.h | 25 +- internal/core/unittest/CMakeLists.txt | 18 +- internal/core/unittest/bench/bench_search.cpp | 10 +- internal/core/unittest/init_gtest.cpp | 3 + .../unittest/test_array_inverted_index.cpp | 5 +- internal/core/unittest/test_binlog_index.cpp | 108 +- internal/core/unittest/test_c_api.cpp | 292 ++--- .../unittest/test_cached_search_iterator.cpp | 27 +- .../unittest/test_cachinglayer/CMakeLists.txt | 29 + .../cachinglayer_test_utils.h | 171 +++ .../test_cachinglayer/mock_list_node.h | 85 ++ .../test_cachinglayer/test_cache_slot.cpp | 846 ++++++++++++++ .../unittest/test_cachinglayer/test_dlist.cpp | 610 ++++++++++ internal/core/unittest/test_chunk.cpp | 76 +- internal/core/unittest/test_chunk_cache.cpp | 251 ---- .../core/unittest/test_chunked_column.cpp | 27 +- .../core/unittest/test_chunked_segment.cpp | 162 ++- internal/core/unittest/test_delete_record.cpp | 29 +- internal/core/unittest/test_exec.cpp | 23 +- internal/core/unittest/test_expr.cpp | 338 +----- .../unittest/test_expr_materialized_view.cpp | 52 +- internal/core/unittest/test_group_by.cpp | 95 +- internal/core/unittest/test_growing.cpp | 11 +- .../core/unittest/test_growing_storage_v2.cpp | 4 +- .../core/unittest/test_inverted_index.cpp | 42 +- .../core/unittest/test_iterative_filter.cpp | 63 +- .../core/unittest/test_kmeans_clustering.cpp | 2 +- internal/core/unittest/test_query.cpp | 10 +- internal/core/unittest/test_random_sample.cpp | 12 +- internal/core/unittest/test_regex_query.cpp | 10 +- internal/core/unittest/test_retrieve.cpp | 28 +- .../unittest/test_scalar_index_creator.cpp | 8 - internal/core/unittest/test_sealed.cpp | 695 ++--------- internal/core/unittest/test_segcore.cpp | 84 -- internal/core/unittest/test_span.cpp | 13 +- internal/core/unittest/test_storage.cpp | 98 +- internal/core/unittest/test_text_match.cpp | 23 +- internal/core/unittest/test_utils.cpp | 3 +- .../core/unittest/test_utils/AssertUtils.h | 3 +- internal/core/unittest/test_utils/Constants.h | 3 + internal/core/unittest/test_utils/DataGen.h | 84 +- .../unittest/test_utils/storage_test_utils.h | 171 ++- .../segments/load_field_data_info.go | 10 - internal/querynodev2/segments/segment.go | 13 - .../querynodev2/segments/segment_loader.go | 2 +- .../segments/segment_loader_test.go | 2 + internal/querynodev2/server.go | 6 + internal/util/segcore/segment.go | 3 - pkg/util/paramtable/component_param.go | 45 + scripts/run_cpp_unittest.sh | 7 + 134 files changed, 6262 insertions(+), 4439 deletions(-) create mode 100644 internal/core/src/cachinglayer/CMakeLists.txt create mode 100644 internal/core/src/cachinglayer/CacheSlot.h create mode 100644 internal/core/src/cachinglayer/Manager.cpp create mode 100644 internal/core/src/cachinglayer/Manager.h create mode 100644 internal/core/src/cachinglayer/Translator.h create mode 100644 internal/core/src/cachinglayer/Utils.h create mode 100644 internal/core/src/cachinglayer/lrucache/DList.cpp create mode 100644 internal/core/src/cachinglayer/lrucache/DList.h create mode 100644 internal/core/src/cachinglayer/lrucache/ListNode.cpp create mode 100644 internal/core/src/cachinglayer/lrucache/ListNode.h create mode 100644 internal/core/src/common/ArrowDataWrapper.h delete mode 100644 internal/core/src/mmap/Column.h create mode 100644 internal/core/src/segcore/storagev1translator/ChunkTranslator.cpp create mode 100644 internal/core/src/segcore/storagev1translator/ChunkTranslator.h create mode 100644 internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.cpp create mode 100644 internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.h create mode 100644 internal/core/src/segcore/storagev1translator/InsertRecordTranslator.cpp create mode 100644 internal/core/src/segcore/storagev1translator/InsertRecordTranslator.h delete mode 100644 internal/core/src/storage/ChunkCache.cpp delete mode 100644 internal/core/src/storage/ChunkCache.h create mode 100644 internal/core/unittest/test_cachinglayer/CMakeLists.txt create mode 100644 internal/core/unittest/test_cachinglayer/cachinglayer_test_utils.h create mode 100644 internal/core/unittest/test_cachinglayer/mock_list_node.h create mode 100644 internal/core/unittest/test_cachinglayer/test_cache_slot.cpp create mode 100644 internal/core/unittest/test_cachinglayer/test_dlist.cpp delete mode 100644 internal/core/unittest/test_chunk_cache.cpp diff --git a/.gitignore b/.gitignore index 3e5bc6c8c5..6f47fb5ec2 100644 --- a/.gitignore +++ b/.gitignore @@ -9,7 +9,9 @@ **/cmake-build-release/* **/cmake_build_release/* **/cmake_build/* -**/CmakeFiles/* +**/CMakeFiles/* +**/.cmake/* +CMakeCache.txt .cache coverage_report/ diff --git a/configs/milvus.yaml b/configs/milvus.yaml index 0de44d8111..43aedbd410 100644 --- a/configs/milvus.yaml +++ b/configs/milvus.yaml @@ -98,7 +98,7 @@ minio: # minio.address and minio.port together generate the valid access to MinIO or S3 service. # MinIO preferentially acquires the valid IP address from the environment variable MINIO_ADDRESS when Milvus is started. # Default value applies when MinIO or S3 is running on the same network with Milvus. - address: localhost + address: localhost:9000 port: 9000 # Port of MinIO or S3 service. # Access key ID that MinIO or S3 issues to user for authorized access. # Environment variable: MINIO_ACCESS_KEY_ID or minio.accessKeyID @@ -447,6 +447,10 @@ queryNode: memExpansionRate: 1.15 # extra memory needed by building interim index buildParallelRate: 0.5 # the ratio of building interim index parallel matched with cpu num multipleChunkedEnable: true # Deprecated. Enable multiple chunked search + tieredStorage: + enableGlobally: true # Whether or not to turn on Tiered Storage globally in this cluster. + memoryAllocationRatio: 0.8 # The ratio of memory allocation for Tiered Storage. + diskAllocationRatio: 0.8 # The ratio of disk allocation for Tiered Storage. knowhereScoreConsistency: false # Enable knowhere strong consistency score computation logic jsonKeyStatsCommitInterval: 200 # the commit interval for the JSON key Stats to commit loadMemoryUsageFactor: 1 # The multiply factor of calculating the memory usage while loading segments diff --git a/docs/design_docs/segcore/segment_sealed.md b/docs/design_docs/segcore/segment_sealed.md index 86960ee020..267e967cdc 100644 --- a/docs/design_docs/segcore/segment_sealed.md +++ b/docs/design_docs/segcore/segment_sealed.md @@ -8,7 +8,6 @@ SegmentSealed has an extra interface rather than SegmentInterface: 2. `LoadFieldData(loadFieldDataInfo)`: load column data, could be either scalar column or vector column 1. Note: indexes and vector data for the same column may coexist. Indexes are prioritized in the search 3. `DropIndex(fieldId)`: drop and release an existing index of a specified field -4. `DropFieldData(fieldId)`: drop and release existing data for a specified field Search is executable as long as all the columns involved in the search are loaded. diff --git a/internal/core/conanfile.py b/internal/core/conanfile.py index de9cb8a345..59a73843da 100644 --- a/internal/core/conanfile.py +++ b/internal/core/conanfile.py @@ -71,7 +71,7 @@ class MilvusConan(ConanFile): "aws-sdk-cpp:config": True, "aws-sdk-cpp:text-to-speech": False, "aws-sdk-cpp:transfer": False, - "gtest:build_gmock": False, + "gtest:build_gmock": True, "boost:without_locale": False, "boost:without_test": True, "glog:with_gflags": True, diff --git a/internal/core/src/CMakeLists.txt b/internal/core/src/CMakeLists.txt index 1b1baa28b2..8e41254d95 100644 --- a/internal/core/src/CMakeLists.txt +++ b/internal/core/src/CMakeLists.txt @@ -49,6 +49,7 @@ add_subdirectory( clustering ) add_subdirectory( exec ) add_subdirectory( bitset ) add_subdirectory( futures ) +add_subdirectory( cachinglayer ) milvus_add_pkg_config("milvus_core") @@ -67,6 +68,7 @@ add_library(milvus_core SHARED $ $ $ + $ ) set(LINK_TARGETS diff --git a/internal/core/src/cachinglayer/CMakeLists.txt b/internal/core/src/cachinglayer/CMakeLists.txt new file mode 100644 index 0000000000..5115a87363 --- /dev/null +++ b/internal/core/src/cachinglayer/CMakeLists.txt @@ -0,0 +1,13 @@ +# Copyright (C) 2019-2025 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under the License + +add_source_at_current_directory_recursively() +add_library(milvus_cachinglayer OBJECT ${SOURCE_FILES}) diff --git a/internal/core/src/cachinglayer/CacheSlot.h b/internal/core/src/cachinglayer/CacheSlot.h new file mode 100644 index 0000000000..fa6c2a8c7d --- /dev/null +++ b/internal/core/src/cachinglayer/CacheSlot.h @@ -0,0 +1,372 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "cachinglayer/lrucache/DList.h" +#include "cachinglayer/lrucache/ListNode.h" +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" +#include "log/Log.h" +#include "monitor/prometheus_client.h" + +namespace milvus::cachinglayer { + +template +class CellAccessor; + +// - The action of pinning cells is not started until the returned SemiFuture is scheduled on an executor. +// - Once the future is scheduled, CacheSlot must live until the future is ready. +// - The returned CellAccessor stores a shared_ptr of CacheSlot, thus will keep CacheSlot alive. +template +class CacheSlot final : public std::enable_shared_from_this> { + public: + // TODO(tiered storage 1): the CellT should return its actual usage, once loaded. And we use this to report metrics. + static_assert( + std::is_same_v().CellByteSize())>, + "CellT must have a CellByteSize() method that returns a size_t " + "representing the memory consumption of the cell"); + + CacheSlot(std::unique_ptr> translator, + internal::DList* dlist) + : translator_(std::move(translator)), + cells_(translator_->num_cells()), + dlist_(dlist) { + for (cid_t i = 0; i < translator_->num_cells(); ++i) { + new (&cells_[i]) + CacheCell(this, i, translator_->estimated_byte_size_of_cell(i)); + } + internal::cache_slot_count(translator_->meta()->storage_type) + .Increment(); + internal::cache_cell_count(translator_->meta()->storage_type) + .Increment(translator_->num_cells()); + } + + CacheSlot(const CacheSlot&) = delete; + CacheSlot& + operator=(const CacheSlot&) = delete; + CacheSlot(CacheSlot&&) = delete; + CacheSlot& + operator=(CacheSlot&&) = delete; + + folly::SemiFuture>> + PinCells(std::vector uids) { + return folly::makeSemiFuture().deferValue([this, + uids = std::vector( + uids)](auto&&) { + auto count = uids.size(); + std::unordered_set involved_cids; + involved_cids.reserve(count); + for (size_t i = 0; i < count; ++i) { + auto uid = uids[i]; + auto cid = translator_->cell_id_of(uid); + if (cid >= cells_.size()) { + return folly::makeSemiFuture< + std::shared_ptr>>( + folly::make_exception_wrapper( + fmt::format( + "CacheSlot {}: translator returned cell_id {} " + "for uid {} which is out of range", + translator_->key(), + cid, + uid))); + } + involved_cids.insert(cid); + } + + std::vector> futures; + std::unordered_set need_load_cids; + futures.reserve(involved_cids.size()); + need_load_cids.reserve(involved_cids.size()); + for (auto cid : involved_cids) { + auto [need_load, future] = cells_[cid].pin(); + futures.push_back(std::move(future)); + if (need_load) { + need_load_cids.insert(cid); + } + } + auto load_future = folly::makeSemiFuture(); + if (!need_load_cids.empty()) { + load_future = RunLoad(std::move(need_load_cids)); + } + return std::move(load_future) + .deferValue( + [this, futures = std::move(futures)](auto&&) mutable + -> folly::SemiFuture>> { + return folly::collect(futures).deferValue( + [this](std::vector&& + pins) mutable + -> std::shared_ptr> { + return std::make_shared>( + this->shared_from_this(), std::move(pins)); + }); + }); + }); + } + + size_t + num_cells() const { + return translator_->num_cells(); + } + + ResourceUsage + size_of_cell(cid_t cid) const { + return translator_->estimated_byte_size_of_cell(cid); + } + + Meta* + meta() { + return translator_->meta(); + } + + ~CacheSlot() { + internal::cache_slot_count(translator_->meta()->storage_type) + .Decrement(); + internal::cache_cell_count(translator_->meta()->storage_type) + .Decrement(translator_->num_cells()); + } + + private: + friend class CellAccessor; + + cid_t + cell_id_of(uid_t uid) const { + return translator_->cell_id_of(uid); + } + + folly::SemiFuture + RunLoad(std::unordered_set&& cids) { + return folly::makeSemiFuture().deferValue( + [this, + cids = std::move(cids)](auto&&) -> folly::SemiFuture { + try { + auto start = std::chrono::high_resolution_clock::now(); + std::vector cids_vec(cids.begin(), cids.end()); + auto results = translator_->get_cells(cids_vec); + auto latency = + std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - start); + for (auto& result : results) { + cells_[result.first].set_cell( + std::move(result.second), + cids.count(result.first) > 0); + internal::cache_load_latency( + translator_->meta()->storage_type) + .Observe(latency.count()); + } + internal::cache_cell_loaded_count( + translator_->meta()->storage_type) + .Increment(results.size()); + internal::cache_load_count_success( + translator_->meta()->storage_type) + .Increment(results.size()); + } catch (...) { + auto exception = std::current_exception(); + auto ew = folly::exception_wrapper(exception); + internal::cache_load_count_fail( + translator_->meta()->storage_type) + .Increment(cids.size()); + for (auto cid : cids) { + cells_[cid].set_error(ew); + } + return folly::makeSemiFuture(ew); + } + return folly::Unit(); + }); + } + + struct CacheCell : internal::ListNode { + public: + CacheCell() = default; + CacheCell(CacheSlot* slot, cid_t cid, ResourceUsage size) + : internal::ListNode(slot->dlist_, size), slot_(slot), cid_(cid) { + } + ~CacheCell() { + if (state_ == State::LOADING) { + LOG_ERROR("CacheSlot Cell {} destroyed while loading", key()); + } + } + + CellT* + cell() { + return cell_.get(); + } + + // Be careful that even though only a single thread can request loading a cell, + // it is still possible that multiple threads call set_cell() concurrently. + // For example, 2 RunLoad() calls tries to download cell 4 and 6, and both decided + // to also download cell 5, if they finished at the same time, they will call set_cell() + // of cell 5 concurrently. + void + set_cell(std::unique_ptr cell, bool requesting_thread) { + mark_loaded( + [this, cell = std::move(cell)]() mutable { + cell_ = std::move(cell); + life_start_ = std::chrono::steady_clock::now(); + milvus::monitor::internal_cache_used_bytes_memory.Increment( + size_.memory_bytes); + milvus::monitor::internal_cache_used_bytes_disk.Increment( + size_.file_bytes); + }, + requesting_thread); + } + + void + set_error(folly::exception_wrapper error) { + internal::ListNode::set_error(std::move(error)); + } + + protected: + void + unload() override { + if (cell_) { + internal::cache_cell_loaded_count( + slot_->translator_->meta()->storage_type) + .Decrement(); + auto life_time = std::chrono::steady_clock::now() - life_start_; + auto seconds = + std::chrono::duration_cast(life_time) + .count(); + internal::cache_item_lifetime_seconds( + slot_->translator_->meta()->storage_type) + .Observe(seconds); + cell_ = nullptr; + milvus::monitor::internal_cache_used_bytes_memory.Decrement( + size_.memory_bytes); + milvus::monitor::internal_cache_used_bytes_disk.Decrement( + size_.file_bytes); + } + } + std::string + key() const override { + return fmt::format("{}:{}", slot_->translator_->key(), cid_); + } + + private: + CacheSlot* slot_{nullptr}; + cid_t cid_{0}; + std::unique_ptr cell_{nullptr}; + std::chrono::steady_clock::time_point life_start_{}; + }; + const std::unique_ptr> translator_; + // Each CacheCell's cid_t is its index in vector + // Once initialized, cells_ should never be resized. + std::vector cells_; + internal::DList* dlist_; +}; + +// - A thin wrapper for accessing cells in a CacheSlot. +// - When this class is created, the cells are loaded and pinned. +// - Accessing cells through this class does not incur any lock overhead. +// - Accessing cells that are not pinned by this CellAccessor is undefined behavior. +template +class CellAccessor { + public: + CellAccessor(std::shared_ptr> slot, + std::vector pins) + : slot_(std::move(slot)), pins_(std::move(pins)) { + } + + CellT* + get_cell_of(uid_t uid) { + auto cid = slot_->cell_id_of(uid); + return slot_->cells_[cid].cell(); + } + + CellT* + get_ith_cell(cid_t cid) { + return slot_->cells_[cid].cell(); + } + + private: + // pins must be destroyed before slot_ is destroyed, thus + // pins_ should be a member after slot_. + std::shared_ptr> slot_; + std::vector pins_; +}; + +// TODO(tiered storage 4): this class is a temp solution. Later we should modify all usage of this class +// to use folly::SemiFuture instead: all data access should happen within deferValue(). +// Current impl requires the T type to be movable/copyable. +template +class PinWrapper { + public: + PinWrapper() = default; + PinWrapper(std::any raii, T&& content) + : raii_(std::move(raii)), content_(std::move(content)) { + } + + PinWrapper(std::any raii, const T& content) + : raii_(std::move(raii)), content_(content) { + } + + // For those that does not need a pin. eg: growing segment, views that actually copies the data, etc. + PinWrapper(T&& content) : raii_(nullptr), content_(std::move(content)) { + } + PinWrapper(const T& content) : raii_(nullptr), content_(content) { + } + + PinWrapper(PinWrapper&& other) noexcept + : raii_(std::move(other.raii_)), content_(std::move(other.content_)) { + } + + PinWrapper(const PinWrapper& other) + : raii_(other.raii_), content_(other.content_) { + } + + PinWrapper& + operator=(PinWrapper&& other) noexcept { + if (this != &other) { + raii_ = std::move(other.raii_); + content_ = std::move(other.content_); + } + return *this; + } + + PinWrapper& + operator=(const PinWrapper& other) { + if (this != &other) { + raii_ = other.raii_; + content_ = other.content_; + } + return *this; + } + + T& + get() { + return content_; + } + + template + PinWrapper + transform(Fn&& transformer) && { + T2 transformed = transformer(std::move(content_)); + return PinWrapper(std::move(raii_), std::move(transformed)); + } + + private: + // CellAccessor is templated on CellT, we don't want to enforce that in this class. + std::any raii_{nullptr}; + T content_; +}; + +} // namespace milvus::cachinglayer diff --git a/internal/core/src/cachinglayer/Manager.cpp b/internal/core/src/cachinglayer/Manager.cpp new file mode 100644 index 0000000000..870d17c8f3 --- /dev/null +++ b/internal/core/src/cachinglayer/Manager.cpp @@ -0,0 +1,62 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#include "cachinglayer/Manager.h" + +#include + +#include "cachinglayer/Utils.h" +#include "log/Log.h" + +namespace milvus::cachinglayer { + +Manager& +Manager::GetInstance() { + static Manager instance; + return instance; +} + +bool +Manager::ConfigureTieredStorage(bool enabled_globally, + int64_t memory_limit_bytes, + int64_t disk_limit_bytes) { + Manager& manager = GetInstance(); + if (enabled_globally) { + if (manager.dlist_ != nullptr) { + return manager.dlist_->UpdateLimit( + {memory_limit_bytes, disk_limit_bytes}); + } else { + ResourceUsage limit{memory_limit_bytes, disk_limit_bytes}; + internal::DList::TouchConfig touch_config{std::chrono::seconds(10)}; + manager.dlist_ = + std::make_unique(limit, touch_config); + } + LOG_INFO( + "Configured Tiered Storage manager with memory limit: {} bytes " + "({:.2f} GB), disk " + "limit: {} bytes ({:.2f} GB)", + memory_limit_bytes, + memory_limit_bytes / (1024.0 * 1024.0 * 1024.0), + disk_limit_bytes, + disk_limit_bytes / (1024.0 * 1024.0 * 1024.0)); + } else { + manager.dlist_ = nullptr; + LOG_INFO("Tiered Storage is disabled"); + } + return true; +} + +size_t +Manager::memory_overhead() const { + // TODO(tiered storage 2): calculate memory overhead + return 0; +} + +} // namespace milvus::cachinglayer diff --git a/internal/core/src/cachinglayer/Manager.h b/internal/core/src/cachinglayer/Manager.h new file mode 100644 index 0000000000..b8e99d1632 --- /dev/null +++ b/internal/core/src/cachinglayer/Manager.h @@ -0,0 +1,66 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include + +#include "cachinglayer/CacheSlot.h" +#include "cachinglayer/Translator.h" +#include "cachinglayer/lrucache/DList.h" + +namespace milvus::cachinglayer { + +class Manager { + public: + static Manager& + GetInstance(); + + // This function is not thread safe, must be called before any CacheSlot is created. + // TODO(tiered storage 4): support dynamic update. + static bool + ConfigureTieredStorage(bool enabled_globally, + int64_t memory_limit_bytes, + int64_t disk_limit_bytes); + + Manager(const Manager&) = delete; + Manager& + operator=(const Manager&) = delete; + Manager(Manager&&) = delete; + Manager& + operator=(Manager&&) = delete; + + ~Manager() = default; + + template + std::shared_ptr> + CreateCacheSlot(std::unique_ptr> translator) { + return std::make_shared>(std::move(translator), + dlist_.get()); + } + + // memory overhead for managing all cache slots/cells/translators/policies. + size_t + memory_overhead() const; + + private: + friend void + ConfigureTieredStorage(bool enabled_globally, + int64_t memory_limit_bytes, + int64_t disk_limit_bytes); + + Manager() = default; // Private constructor + + std::unique_ptr dlist_{nullptr}; + bool enable_global_tiered_storage_{false}; +}; // class Manager + +} // namespace milvus::cachinglayer diff --git a/internal/core/src/cachinglayer/Translator.h b/internal/core/src/cachinglayer/Translator.h new file mode 100644 index 0000000000..7f703a040d --- /dev/null +++ b/internal/core/src/cachinglayer/Translator.h @@ -0,0 +1,60 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include +#include + +#include "cachinglayer/Utils.h" + +namespace milvus::cachinglayer { + +struct Meta { + StorageType storage_type; + Meta(StorageType storage_type) : storage_type(storage_type) { + } +}; + +template +class Translator { + public: + using value_type = CellT; + + virtual size_t + num_cells() const = 0; + virtual cid_t + cell_id_of(uid_t uid) const = 0; + // For resource reservation when a cell is about to be loaded. + // If a cell is about to be pinned and loaded, and there are not enough resource for it, EvictionManager + // will try to evict some other cells to make space. Thus this estimation should generally be greater + // than or equal to the actual size. If the estimation is smaller than the actual size, with insufficient + // resource reserved, the load may fail. + virtual ResourceUsage + estimated_byte_size_of_cell(cid_t cid) const = 0; + // must be unique to identify a CacheSlot. + virtual const std::string& + key() const = 0; + + virtual Meta* + meta() = 0; + + // Translator may choose to fetch more than requested cells. + // TODO(tiered storage 2): This has a problem: when loading, the resource manager will only reserve the size of the + // requested cells, How can translator be sure the extra cells can fit? Currently if bonus cells are returned, + // used memory in cache may exceed the limit. Maybe try to reserve memory for bonus cells, and drop cell if failed. + virtual std::vector>> + get_cells(const std::vector& cids) = 0; + virtual ~Translator() = default; +}; + +} // namespace milvus::cachinglayer diff --git a/internal/core/src/cachinglayer/Utils.h b/internal/core/src/cachinglayer/Utils.h new file mode 100644 index 0000000000..4fe1833d5e --- /dev/null +++ b/internal/core/src/cachinglayer/Utils.h @@ -0,0 +1,300 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "common/EasyAssert.h" +#include "folly/executors/InlineExecutor.h" +#include +#include +#include +#include + +#include "monitor/prometheus_client.h" + +namespace milvus::cachinglayer { + +using uid_t = int64_t; +using cid_t = int64_t; + +enum class StorageType { + MEMORY, + DISK, + MIXED, +}; + +// TODO(tiered storage 4): this is a temporary function to get the result of a future +// by running it on the inline executor. We don't need this once we are fully async. +template +T +SemiInlineGet(folly::SemiFuture&& future) { + return std::move(future).via(&folly::InlineExecutor::instance()).get(); +} + +struct ResourceUsage { + int64_t memory_bytes{0}; + int64_t file_bytes{0}; + + ResourceUsage() noexcept : memory_bytes(0), file_bytes(0) { + } + ResourceUsage(int64_t mem, int64_t file) noexcept + : memory_bytes(mem), file_bytes(file) { + } + + ResourceUsage + operator+(const ResourceUsage& rhs) const { + return ResourceUsage(memory_bytes + rhs.memory_bytes, + file_bytes + rhs.file_bytes); + } + + void + operator+=(const ResourceUsage& rhs) { + memory_bytes += rhs.memory_bytes; + file_bytes += rhs.file_bytes; + } + + ResourceUsage + operator-(const ResourceUsage& rhs) const { + return ResourceUsage(memory_bytes - rhs.memory_bytes, + file_bytes - rhs.file_bytes); + } + + void + operator-=(const ResourceUsage& rhs) { + memory_bytes -= rhs.memory_bytes; + file_bytes -= rhs.file_bytes; + } + + bool + operator==(const ResourceUsage& rhs) const { + return memory_bytes == rhs.memory_bytes && file_bytes == rhs.file_bytes; + } + + bool + operator!=(const ResourceUsage& rhs) const { + return !(*this == rhs); + } + + bool + GEZero() const { + return memory_bytes >= 0 && file_bytes >= 0; + } + + bool + CanHold(const ResourceUsage& rhs) const { + return memory_bytes >= rhs.memory_bytes && file_bytes >= rhs.file_bytes; + } + + StorageType + storage_type() const { + if (memory_bytes > 0 && file_bytes > 0) { + return StorageType::MIXED; + } + return memory_bytes > 0 ? StorageType::MEMORY : StorageType::DISK; + } + + std::string + ToString() const { + return fmt::format("ResourceUsage{memory_bytes={}, file_bytes={}}", + memory_bytes, + file_bytes); + } +}; + +inline std::ostream& +operator<<(std::ostream& os, const ResourceUsage& usage) { + os << "ResourceUsage{memory_bytes=" << usage.memory_bytes + << ", file_bytes=" << usage.file_bytes << "}"; + return os; +} + +inline void +operator+=(std::atomic& atomic_lhs, const ResourceUsage& rhs) { + ResourceUsage current = atomic_lhs.load(); + ResourceUsage new_value; + do { + new_value = current; + new_value += rhs; + } while (!atomic_lhs.compare_exchange_weak(current, new_value)); +} + +inline void +operator-=(std::atomic& atomic_lhs, const ResourceUsage& rhs) { + ResourceUsage current = atomic_lhs.load(); + ResourceUsage new_value; + do { + new_value = current; + new_value -= rhs; + } while (!atomic_lhs.compare_exchange_weak(current, new_value)); +} + +namespace internal { + +inline prometheus::Gauge& +cache_slot_count(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_slot_count_memory; + case StorageType::DISK: + return monitor::internal_cache_slot_count_disk; + case StorageType::MIXED: + return monitor::internal_cache_slot_count_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Gauge& +cache_cell_count(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_cell_count_memory; + case StorageType::DISK: + return monitor::internal_cache_cell_count_disk; + case StorageType::MIXED: + return monitor::internal_cache_cell_count_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Gauge& +cache_cell_loaded_count(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_cell_loaded_count_memory; + case StorageType::DISK: + return monitor::internal_cache_cell_loaded_count_disk; + case StorageType::MIXED: + return monitor::internal_cache_cell_loaded_count_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Histogram& +cache_load_latency(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_load_latency_memory; + case StorageType::DISK: + return monitor::internal_cache_load_latency_disk; + case StorageType::MIXED: + return monitor::internal_cache_load_latency_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Counter& +cache_op_result_count_hit(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_op_result_count_hit_memory; + case StorageType::DISK: + return monitor::internal_cache_op_result_count_hit_disk; + case StorageType::MIXED: + return monitor::internal_cache_op_result_count_hit_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Counter& +cache_op_result_count_miss(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_op_result_count_miss_memory; + case StorageType::DISK: + return monitor::internal_cache_op_result_count_miss_disk; + case StorageType::MIXED: + return monitor::internal_cache_op_result_count_miss_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Counter& +cache_eviction_count(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_eviction_count_memory; + case StorageType::DISK: + return monitor::internal_cache_eviction_count_disk; + case StorageType::MIXED: + return monitor::internal_cache_eviction_count_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Histogram& +cache_item_lifetime_seconds(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_item_lifetime_seconds_memory; + case StorageType::DISK: + return monitor::internal_cache_item_lifetime_seconds_disk; + case StorageType::MIXED: + return monitor::internal_cache_item_lifetime_seconds_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Counter& +cache_load_count_success(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_load_count_success_memory; + case StorageType::DISK: + return monitor::internal_cache_load_count_success_disk; + case StorageType::MIXED: + return monitor::internal_cache_load_count_success_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Counter& +cache_load_count_fail(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_load_count_fail_memory; + case StorageType::DISK: + return monitor::internal_cache_load_count_fail_disk; + case StorageType::MIXED: + return monitor::internal_cache_load_count_fail_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +inline prometheus::Gauge& +cache_memory_overhead_bytes(StorageType storage_type) { + switch (storage_type) { + case StorageType::MEMORY: + return monitor::internal_cache_memory_overhead_bytes_memory; + case StorageType::DISK: + return monitor::internal_cache_memory_overhead_bytes_disk; + case StorageType::MIXED: + return monitor::internal_cache_memory_overhead_bytes_mixed; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } +} + +} // namespace internal + +} // namespace milvus::cachinglayer diff --git a/internal/core/src/cachinglayer/lrucache/DList.cpp b/internal/core/src/cachinglayer/lrucache/DList.cpp new file mode 100644 index 0000000000..a4706134f7 --- /dev/null +++ b/internal/core/src/cachinglayer/lrucache/DList.cpp @@ -0,0 +1,201 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#include "cachinglayer/lrucache/DList.h" + +#include +#include + +#include +#include + +#include "cachinglayer/Utils.h" +#include "cachinglayer/lrucache/ListNode.h" +#include "monitor/prometheus_client.h" + +namespace milvus::cachinglayer::internal { + +bool +DList::reserveMemory(const ResourceUsage& size) { + std::unique_lock list_lock(list_mtx_); + auto used = used_memory_.load(); + if (max_memory_.CanHold(used + size)) { + used_memory_ += size; + return true; + } + if (tryEvict(used + size - max_memory_)) { + used_memory_ += size; + return true; + } + return false; +} + +bool +DList::tryEvict(const ResourceUsage& expected_eviction) { + std::vector to_evict; + // items are evicted because they are not used for a while, thus it should be ok to lock them + // a little bit longer. + std::vector> item_locks; + + ResourceUsage size_to_evict; + + auto would_help = [&](const ResourceUsage& size) -> bool { + auto need_memory = + size_to_evict.memory_bytes < expected_eviction.memory_bytes; + auto need_disk = + size_to_evict.file_bytes < expected_eviction.file_bytes; + return (need_memory && size.memory_bytes > 0) || + (need_disk && size.file_bytes > 0); + }; + + for (auto it = tail_; it != nullptr; it = it->next_) { + if (!would_help(it->size())) { + continue; + } + // use try_to_lock to avoid dead lock by failing immediately if the ListNode lock is already held. + auto& lock = item_locks.emplace_back(it->mtx_, std::try_to_lock); + // if lock failed, it means this ListNode will be used again, so we don't evict it anymore. + if (lock.owns_lock() && it->pin_count_ == 0) { + to_evict.push_back(it); + size_to_evict += it->size(); + if (size_to_evict.CanHold(expected_eviction)) { + break; + } + } else { + // if we grabbed the lock only to find that the ListNode is pinned; or if we failed to lock + // the ListNode, we do not evict this ListNode. + item_locks.pop_back(); + } + } + if (!size_to_evict.CanHold(expected_eviction)) { + return false; + } + for (auto* list_node : to_evict) { + auto size = list_node->size(); + internal::cache_eviction_count(size.storage_type()).Increment(); + popItem(list_node); + list_node->clear_data(); + used_memory_ -= size; + } + + switch (size_to_evict.storage_type()) { + case StorageType::MEMORY: + milvus::monitor::internal_cache_evicted_bytes_memory.Increment( + size_to_evict.memory_bytes); + break; + case StorageType::DISK: + milvus::monitor::internal_cache_evicted_bytes_disk.Increment( + size_to_evict.file_bytes); + break; + case StorageType::MIXED: + milvus::monitor::internal_cache_evicted_bytes_memory.Increment( + size_to_evict.memory_bytes); + milvus::monitor::internal_cache_evicted_bytes_disk.Increment( + size_to_evict.file_bytes); + break; + default: + PanicInfo(ErrorCode::UnexpectedError, "Unknown StorageType"); + } + return true; +} + +bool +DList::UpdateLimit(const ResourceUsage& new_limit) { + if (!new_limit.GEZero()) { + throw std::invalid_argument( + "Milvus Caching Layer: memory and disk usage limit must be greater " + "than 0"); + } + std::unique_lock list_lock(list_mtx_); + auto used = used_memory_.load(); + if (!new_limit.CanHold(used)) { + // positive means amount owed + auto deficit = used - new_limit; + if (!tryEvict(deficit)) { + return false; + } + } + max_memory_ = new_limit; + milvus::monitor::internal_cache_capacity_bytes_memory.Set( + max_memory_.memory_bytes); + milvus::monitor::internal_cache_capacity_bytes_disk.Set( + max_memory_.file_bytes); + return true; +} + +void +DList::releaseMemory(const ResourceUsage& size) { + // safe to substract on atomic without lock + used_memory_ -= size; +} + +void +DList::touchItem(ListNode* list_node, std::optional size) { + std::lock_guard list_lock(list_mtx_); + popItem(list_node); + pushHead(list_node); + if (size.has_value()) { + used_memory_ += size.value(); + } +} + +void +DList::removeItem(ListNode* list_node, ResourceUsage size) { + std::lock_guard list_lock(list_mtx_); + if (popItem(list_node)) { + used_memory_ -= size; + } +} + +void +DList::pushHead(ListNode* list_node) { + if (head_ == nullptr) { + head_ = list_node; + tail_ = list_node; + } else { + list_node->prev_ = head_; + head_->next_ = list_node; + head_ = list_node; + } +} + +bool +DList::popItem(ListNode* list_node) { + if (list_node->prev_ == nullptr && list_node->next_ == nullptr && + list_node != head_) { + // list_node is not in the list + return false; + } + if (head_ == tail_) { + head_ = tail_ = nullptr; + list_node->prev_ = list_node->next_ = nullptr; + } else if (head_ == list_node) { + head_ = list_node->prev_; + head_->next_ = nullptr; + list_node->prev_ = nullptr; + } else if (tail_ == list_node) { + tail_ = list_node->next_; + tail_->prev_ = nullptr; + list_node->next_ = nullptr; + } else { + list_node->prev_->next_ = list_node->next_; + list_node->next_->prev_ = list_node->prev_; + list_node->prev_ = list_node->next_ = nullptr; + } + return true; +} + +bool +DList::IsEmpty() const { + std::lock_guard list_lock(list_mtx_); + return head_ == nullptr; +} + +} // namespace milvus::cachinglayer::internal diff --git a/internal/core/src/cachinglayer/lrucache/DList.h b/internal/core/src/cachinglayer/lrucache/DList.h new file mode 100644 index 0000000000..5bcb56d7af --- /dev/null +++ b/internal/core/src/cachinglayer/lrucache/DList.h @@ -0,0 +1,108 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#pragma once + +#include +#include + +#include +#include + +#include "cachinglayer/lrucache/ListNode.h" + +namespace milvus::cachinglayer::internal { + +class DList { + public: + // Touch a node means to move it to the head of the list, which requires locking the entire list. + // Use TouchConfig to reduce the frequency of touching and reduce contention. + struct TouchConfig { + std::chrono::seconds refresh_window = std::chrono::seconds(10); + }; + + DList(ResourceUsage max_memory, TouchConfig touch_config) + : max_memory_(max_memory), touch_config_(touch_config) { + } + + // If after evicting all unpinned items, the used_memory_ is still larger than new_limit, false will be returned + // and no eviction will be done. + // Will throw if new_limit is negative. + bool + UpdateLimit(const ResourceUsage& new_limit); + + // True if no nodes in the list. + bool + IsEmpty() const; + + // This method uses a global lock. + bool + reserveMemory(const ResourceUsage& size); + + // Used only when load failed. This will only cause used_memory_ to decrease, which will not affect the correctness + // of concurrent reserveMemory() even without lock. + void + releaseMemory(const ResourceUsage& size); + + // Caller must guarantee that the current thread holds the lock of list_node->mtx_. + // touchItem is used in 2 places: + // 1. when a loaded cell is pinned/unpinned, we need to touch it to refresh the LRU order. + // we don't update used_memory_ here. + // 2. when a cell is loaded as a bonus, we need to touch it to insert into the LRU and update + // used_memory_ to track the memory usage(usage of such cell is not counted during reservation). + void + touchItem(ListNode* list_node, + std::optional size = std::nullopt); + + // Caller must guarantee that the current thread holds the lock of list_node->mtx_. + // Removes the node from the list and updates used_memory_. + void + removeItem(ListNode* list_node, ResourceUsage size); + + const TouchConfig& + touch_config() const { + return touch_config_; + } + + private: + friend class DListTestFriend; + + // Try to evict some items so that the evicted items are larger than expected_eviction. + // If we cannot achieve the goal, nothing will be evicted and false will be returned. + // Must be called under the lock of list_mtx_. + bool + tryEvict(const ResourceUsage& expected_eviction); + + // Must be called under the lock of list_mtx_ and list_node->mtx_. + // ListNode is guaranteed to be not in the list. + void + pushHead(ListNode* list_node); + + // Must be called under the lock of list_mtx_ and list_node->mtx_. + // If ListNode is not in the list, this function does nothing. + // Returns true if ListNode is in the list and popped, false otherwise. + bool + popItem(ListNode* list_node); + + // head_ is the most recently used item, tail_ is the least recently used item. + // tail_ -> next -> ... -> head_ + // tail_ <- prev <- ... <- head_ + ListNode* head_ = nullptr; + ListNode* tail_ = nullptr; + + // TODO(tiered storage 3): benchmark folly::DistributedMutex for this usecase. + mutable std::mutex list_mtx_; + // access to used_memory_ and max_memory_ must be done under the lock of list_mtx_ + std::atomic used_memory_{}; + ResourceUsage max_memory_; + const TouchConfig touch_config_; +}; + +} // namespace milvus::cachinglayer::internal diff --git a/internal/core/src/cachinglayer/lrucache/ListNode.cpp b/internal/core/src/cachinglayer/lrucache/ListNode.cpp new file mode 100644 index 0000000000..eacd2e5678 --- /dev/null +++ b/internal/core/src/cachinglayer/lrucache/ListNode.cpp @@ -0,0 +1,211 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#include "cachinglayer/lrucache/ListNode.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "cachinglayer/lrucache/DList.h" +#include "cachinglayer/Utils.h" +#include "common/EasyAssert.h" + +namespace milvus::cachinglayer::internal { + +ListNode::NodePin::NodePin(ListNode* node) : node_(node) { + if (node_) { + node_->pin_count_++; + } +} + +ListNode::NodePin::~NodePin() { + if (node_) { + node_->unpin(); + } +} + +ListNode::NodePin::NodePin(NodePin&& other) : NodePin(nullptr) { + std::swap(node_, other.node_); +} + +ListNode::NodePin& +ListNode::NodePin::operator=(NodePin&& other) { + std::swap(node_, other.node_); + return *this; +} + +ListNode::ListNode(DList* dlist, ResourceUsage size) + : last_touch_(dlist ? (std::chrono::high_resolution_clock::now() - + 2 * dlist->touch_config().refresh_window) + : std::chrono::high_resolution_clock::now()), + dlist_(dlist), + size_(size), + state_(State::NOT_LOADED) { +} + +ListNode::~ListNode() { + if (dlist_) { + std::unique_lock lock(mtx_); + dlist_->removeItem(this, size_); + } +} + +ResourceUsage& +ListNode::size() { + return size_; +} + +std::pair> +ListNode::pin() { + // must be called with lock acquired, and state must not be NOT_LOADED. + auto read_op = [this]() -> std::pair> { + AssertInfo(state_ != State::NOT_LOADED, + "Programming error: read_op called on a {} cell", + state_to_string(state_)); + if (state_ == State::ERROR) { + return std::make_pair(false, + folly::makeSemiFuture(error_)); + } + // pin the cell now so that we can avoid taking the lock again in deferValue. + auto p = NodePin(this); + if (state_ == State::LOADED) { + internal::cache_op_result_count_hit(size_.storage_type()) + .Increment(); + return std::make_pair(false, std::move(p)); + } + internal::cache_op_result_count_miss(size_.storage_type()).Increment(); + return std::make_pair(false, + load_promise_->getSemiFuture().deferValue( + [this, p = std::move(p)](auto&&) mutable { + return std::move(p); + })); + }; + { + std::shared_lock lock(mtx_); + if (state_ != State::NOT_LOADED) { + return read_op(); + } + } + std::unique_lock lock(mtx_); + if (state_ != State::NOT_LOADED) { + return read_op(); + } + // need to load. + internal::cache_op_result_count_miss(size_.storage_type()).Increment(); + load_promise_ = std::make_unique>(); + state_ = State::LOADING; + if (dlist_ && !dlist_->reserveMemory(size())) { + // if another thread sees LOADING status, the memory reservation has succeeded. + state_ = State::ERROR; + error_ = folly::make_exception_wrapper(fmt::format( + "Failed to load {} due to insufficient resource", key())); + load_promise_->setException(error_); + load_promise_ = nullptr; + return std::make_pair(false, folly::makeSemiFuture(error_)); + } + + // pin the cell now so that we can avoid taking the lock again in deferValue. + auto p = NodePin(this); + return std::make_pair( + true, + load_promise_->getSemiFuture().deferValue( + [this, p = std::move(p)](auto&&) mutable { return std::move(p); })); +} + +void +ListNode::set_error(folly::exception_wrapper error) { + std::unique_lock lock(mtx_); + AssertInfo(state_ != State::NOT_LOADED && state_ != State::ERROR, + "Programming error: set_error() called on a {} cell", + state_to_string(state_)); + // load failed, release the memory reservation. + if (dlist_) { + dlist_->releaseMemory(size()); + } + // may be successfully loaded by another thread as a bonus, they will update used memory. + if (state_ == State::LOADED) { + return; + } + // else: state_ is LOADING + state_ = State::ERROR; + load_promise_->setException(error); + load_promise_ = nullptr; + error_ = std::move(error); +} + +std::string +ListNode::state_to_string(State state) { + switch (state) { + case State::NOT_LOADED: + return "NOT_LOADED"; + case State::LOADING: + return "LOADING"; + case State::LOADED: + return "LOADED"; + case State::ERROR: + return "ERROR"; + } + throw std::invalid_argument("Invalid state"); +} + +void +ListNode::unpin() { + std::unique_lock lock(mtx_); + AssertInfo( + state_ == State::LOADED || state_ == State::ERROR, + "Programming error: unpin() called on a {} cell, current pin_count {}", + state_to_string(state_), + pin_count_.load()); + if (state_ == State::ERROR) { + return; + } + if (pin_count_.fetch_sub(1) == 1) { + touch(false); + } +} + +void +ListNode::touch(bool update_used_memory) { + auto now = std::chrono::high_resolution_clock::now(); + if (dlist_ && now - last_touch_ > dlist_->touch_config().refresh_window) { + std::optional size = std::nullopt; + if (update_used_memory) { + size = size_; + } + dlist_->touchItem(this, size); + last_touch_ = now; + } +} + +void +ListNode::clear_data() { + // if the cell is evicted, loaded, pinned and unpinned within a single refresh window, + // the cell should be inserted into the cache again. + if (dlist_) { + last_touch_ = std::chrono::high_resolution_clock::now() - + 2 * dlist_->touch_config().refresh_window; + } + unload(); + state_ = State::NOT_LOADED; +} + +void +ListNode::unload() { + // Default implementation does nothing +} + +} // namespace milvus::cachinglayer::internal diff --git a/internal/core/src/cachinglayer/lrucache/ListNode.h b/internal/core/src/cachinglayer/lrucache/ListNode.h new file mode 100644 index 0000000000..95f757b47f --- /dev/null +++ b/internal/core/src/cachinglayer/lrucache/ListNode.h @@ -0,0 +1,183 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include "cachinglayer/Utils.h" +#include "common/EasyAssert.h" + +namespace milvus::cachinglayer::internal { + +class DList; + +// ListNode is not movable/copyable. +class ListNode { + public: + // RAII class to unpin the node. + class NodePin { + public: + // NodePin is movable but not copyable. + NodePin(NodePin&&); + NodePin& + operator=(NodePin&&); + NodePin(const NodePin&) = delete; + NodePin& + operator=(const NodePin&) = delete; + ~NodePin(); + + private: + NodePin(ListNode* node); + friend class ListNode; + ListNode* node_; + }; + ListNode() = default; + ListNode(DList* dlist, ResourceUsage size); + virtual ~ListNode(); + + // ListNode is not movable/copyable because it contains a shared_mutex. + // ListNode also should not be movable/copyable because that would make + // all NodePin::node_ dangling pointers. + ListNode(const ListNode&) = delete; + ListNode& + operator=(const ListNode&) = delete; + ListNode(ListNode&&) = delete; + ListNode& + operator=(ListNode&&) = delete; + + // bool in return value: whether the caller needs to load this cell. + // - If the cell is already loaded, return false and an immediately ready future with a NodePin, the node is pinned + // upon return. + // - If the cell is in error state, return false and an immediately ready future with an exception. + // - If the cell is already being loaded by another thread, return false and a future that will be ready when the + // cell is loaded. The node will not be pinned until the future is ready. + // - Otherwise, the cell is not loaded and not being loaded, return true and a future that will be ready when the + // cell is loaded. The caller needs to load this cell and call mark_loaded() to set the cell as loaded. + // The node will not be pinned until the future is ready. + std::pair> + pin(); + + ResourceUsage& + size(); + + // TODO(tiered storage 1): pin on ERROR should re-trigger loading. + // NOT_LOADED ---> LOADING ---> ERROR + // ^ | + // | v + // |------- LOADED + enum class State { NOT_LOADED, LOADING, LOADED, ERROR }; + + protected: + // will be called during eviction, implementation should release all resources. + virtual void + unload(); + + virtual std::string + key() const = 0; + + template + void + mark_loaded(Fn&& cb, bool requesting_thread) { + std::unique_lock lock(mtx_); + if (requesting_thread) { + // requesting thread will promote NOT_LOADED to LOADING and only requesting + // thread will set state to ERROR, thus it is not possible for the requesting + // thread to see NOT_LOADED or ERROR. + AssertInfo(state_ != State::NOT_LOADED && state_ != State::ERROR, + "Programming error: mark_loaded(requesting_thread=true) " + "called on a {} cell", + state_to_string(state_)); + // no need to touch() here: node is pinned thus not eligible for eviction. + // we can delay touch() to when unpin() is called. + if (state_ == State::LOADING) { + cb(); + state_ = State::LOADED; + load_promise_->setValue(folly::Unit()); + load_promise_ = nullptr; + } else { + // LOADED: cell has been loaded by another thread, do nothing. + return; + } + } else { + // Even though this thread did not request loading this cell, translator still + // decided to download it because the adjacent cells are requested. + if (state_ == State::NOT_LOADED || state_ == State::ERROR) { + state_ = State::LOADED; + cb(); + // memory of this cell is not reserved, touch() to track it. + touch(true); + } else if (state_ == State::LOADING) { + // another thread has explicitly requested loading this cell, we did it first + // thus we set up the state first. + state_ = State::LOADED; + load_promise_->setValue(folly::Unit()); + load_promise_ = nullptr; + cb(); + // the node that marked LOADING has already reserved memory, do not double count. + touch(false); + } else { + // LOADED: cell has been loaded by another thread, do nothing. + return; + } + } + } + + void + set_error(folly::exception_wrapper error); + + State state_{State::NOT_LOADED}; + + static std::string + state_to_string(State state); + + ResourceUsage size_{}; + + private: + friend class DList; + friend class NodePin; + + friend class MockListNode; + friend class DListTest; + friend class DListTestFriend; + friend class ListNodeTestFriend; + friend class ListNodeTest; + + // called by DList during eviction. must be called under the lock of mtx_. + // Made virtual for mock testing. + virtual void + clear_data(); + + void + unpin(); + + // must be called under the lock of mtx_. + void + touch(bool update_used_memory = true); + + mutable std::shared_mutex mtx_; + std::chrono::high_resolution_clock::time_point last_touch_; + // a nullptr dlist_ means this node is not in any DList, and is not prone to cache management. + DList* dlist_; + ListNode* prev_ = nullptr; + ListNode* next_ = nullptr; + std::atomic pin_count_{0}; + + std::unique_ptr> load_promise_{nullptr}; + folly::exception_wrapper error_; +}; + +} // namespace milvus::cachinglayer::internal diff --git a/internal/core/src/common/ArrowDataWrapper.h b/internal/core/src/common/ArrowDataWrapper.h new file mode 100644 index 0000000000..4fdc0a779c --- /dev/null +++ b/internal/core/src/common/ArrowDataWrapper.h @@ -0,0 +1,39 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include + +#include "common/Channel.h" +#include "parquet/arrow/reader.h" + +namespace milvus { + +struct ArrowDataWrapper { + ArrowDataWrapper() = default; + ArrowDataWrapper(std::shared_ptr reader, + std::shared_ptr arrow_reader, + std::shared_ptr file_data) + : reader(std::move(reader)), + arrow_reader(std::move(arrow_reader)), + file_data(std::move(file_data)) { + } + std::shared_ptr reader; + // file reader must outlive the record batch reader + std::shared_ptr arrow_reader; + // underlying file data memory, must outlive the arrow reader + std::shared_ptr file_data; + std::vector> arrow_tables; +}; +using ArrowReaderChannel = Channel>; + +} // namespace milvus diff --git a/internal/core/src/common/Channel.h b/internal/core/src/common/Channel.h index 4a239649b8..fddba45f0d 100644 --- a/internal/core/src/common/Channel.h +++ b/internal/core/src/common/Channel.h @@ -8,13 +8,12 @@ // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License +#pragma once #include -#include #include #include -#include "Exception.h" namespace milvus { template diff --git a/internal/core/src/common/Chunk.h b/internal/core/src/common/Chunk.h index e1f060d488..0efb286841 100644 --- a/internal/core/src/common/Chunk.h +++ b/internal/core/src/common/Chunk.h @@ -11,6 +11,7 @@ #pragma once +#include #include #include #include @@ -54,6 +55,11 @@ class Chunk { return size_; } + size_t + CellByteSize() const { + return size_; + } + int64_t RowNums() const { return row_nums_; diff --git a/internal/core/src/common/ChunkWriter.cpp b/internal/core/src/common/ChunkWriter.cpp index 5b4fac6a7d..00fc488bbb 100644 --- a/internal/core/src/common/ChunkWriter.cpp +++ b/internal/core/src/common/ChunkWriter.cpp @@ -76,14 +76,14 @@ StringChunkWriter::write(const arrow::ArrayVector& array_vec) { } } -std::shared_ptr +std::unique_ptr StringChunkWriter::finish() { // write padding, maybe not needed anymore // FIXME char padding[MMAP_STRING_PADDING]; target_->write(padding, MMAP_STRING_PADDING); auto [data, size] = target_->get(); - return std::make_shared(row_nums_, data, size, nullable_); + return std::make_unique(row_nums_, data, size, nullable_); } void @@ -135,13 +135,13 @@ JSONChunkWriter::write(const arrow::ArrayVector& array_vec) { } } -std::shared_ptr +std::unique_ptr JSONChunkWriter::finish() { char padding[simdjson::SIMDJSON_PADDING]; target_->write(padding, simdjson::SIMDJSON_PADDING); auto [data, size] = target_->get(); - return std::make_shared(row_nums_, data, size, nullable_); + return std::make_unique(row_nums_, data, size, nullable_); } void @@ -219,13 +219,13 @@ ArrayChunkWriter::write(const arrow::ArrayVector& array_vec) { } } -std::shared_ptr +std::unique_ptr ArrayChunkWriter::finish() { char padding[MMAP_ARRAY_PADDING]; target_->write(padding, MMAP_ARRAY_PADDING); auto [data, size] = target_->get(); - return std::make_shared( + return std::make_unique( row_nums_, data, size, element_type_, nullable_); } @@ -283,14 +283,14 @@ SparseFloatVectorChunkWriter::write(const arrow::ArrayVector& array_vec) { } } -std::shared_ptr +std::unique_ptr SparseFloatVectorChunkWriter::finish() { auto [data, size] = target_->get(); - return std::make_shared( + return std::make_unique( row_nums_, data, size, nullable_); } -std::shared_ptr +std::unique_ptr create_chunk(const FieldMeta& field_meta, int dim, const arrow::ArrayVector& array_vec) { @@ -390,7 +390,7 @@ create_chunk(const FieldMeta& field_meta, return w->finish(); } -std::shared_ptr +std::unique_ptr create_chunk(const FieldMeta& field_meta, int dim, File& file, diff --git a/internal/core/src/common/ChunkWriter.h b/internal/core/src/common/ChunkWriter.h index 2fe1552229..8d767ea087 100644 --- a/internal/core/src/common/ChunkWriter.h +++ b/internal/core/src/common/ChunkWriter.h @@ -35,7 +35,7 @@ class ChunkWriterBase { virtual void write(const arrow::ArrayVector& data) = 0; - virtual std::shared_ptr + virtual std::unique_ptr finish() = 0; std::pair @@ -121,10 +121,10 @@ class ChunkWriter final : public ChunkWriterBase { } } - std::shared_ptr + std::unique_ptr finish() override { auto [data, size] = target_->get(); - return std::make_shared( + return std::make_unique( row_nums_, dim_, data, size, sizeof(T), nullable_); } @@ -182,7 +182,7 @@ class StringChunkWriter : public ChunkWriterBase { void write(const arrow::ArrayVector& array_vec) override; - std::shared_ptr + std::unique_ptr finish() override; }; @@ -193,7 +193,7 @@ class JSONChunkWriter : public ChunkWriterBase { void write(const arrow::ArrayVector& array_vec) override; - std::shared_ptr + std::unique_ptr finish() override; }; @@ -212,7 +212,7 @@ class ArrayChunkWriter : public ChunkWriterBase { void write(const arrow::ArrayVector& array_vec) override; - std::shared_ptr + std::unique_ptr finish() override; private: @@ -226,16 +226,16 @@ class SparseFloatVectorChunkWriter : public ChunkWriterBase { void write(const arrow::ArrayVector& array_vec) override; - std::shared_ptr + std::unique_ptr finish() override; }; -std::shared_ptr +std::unique_ptr create_chunk(const FieldMeta& field_meta, int dim, const arrow::ArrayVector& array_vec); -std::shared_ptr +std::unique_ptr create_chunk(const FieldMeta& field_meta, int dim, File& file, @@ -245,4 +245,4 @@ create_chunk(const FieldMeta& field_meta, arrow::ArrayVector read_single_column_batches(std::shared_ptr reader); -} // namespace milvus \ No newline at end of file +} // namespace milvus diff --git a/internal/core/src/common/FieldData.h b/internal/core/src/common/FieldData.h index 873db7af92..7341953ac7 100644 --- a/internal/core/src/common/FieldData.h +++ b/internal/core/src/common/FieldData.h @@ -24,7 +24,7 @@ #include "common/FieldDataInterface.h" #include "common/Channel.h" -#include "parquet/arrow/reader.h" +#include "common/ArrowDataWrapper.h" namespace milvus { @@ -156,25 +156,7 @@ using FieldDataPtr = std::shared_ptr; using FieldDataChannel = Channel; using FieldDataChannelPtr = std::shared_ptr; -struct ArrowDataWrapper { - ArrowDataWrapper() = default; - ArrowDataWrapper(std::shared_ptr reader, - std::shared_ptr arrow_reader, - std::shared_ptr file_data) - : reader(std::move(reader)), - arrow_reader(std::move(arrow_reader)), - file_data(std::move(file_data)) { - } - std::shared_ptr reader; - // file reader must outlive the record batch reader - std::shared_ptr arrow_reader; - // underlying file data memory, must outlive the arrow reader - std::shared_ptr file_data; - std::vector> arrow_tables; -}; -using ArrowReaderChannel = Channel>; - FieldDataPtr InitScalarFieldData(const DataType& type, bool nullable, int64_t cap_rows); -} // namespace milvus \ No newline at end of file +} // namespace milvus diff --git a/internal/core/src/common/IndexMeta.h b/internal/core/src/common/IndexMeta.h index 132c1fd34b..58460aa379 100644 --- a/internal/core/src/common/IndexMeta.h +++ b/internal/core/src/common/IndexMeta.h @@ -95,4 +95,8 @@ class CollectionIndexMeta { using IndexMetaPtr = std::shared_ptr; +const static IndexMetaPtr empty_index_meta = + std::make_shared(1024, + std::map()); + } //namespace milvus \ No newline at end of file diff --git a/internal/core/src/common/LoadInfo.h b/internal/core/src/common/LoadInfo.h index 00774b8162..935c2b43ed 100644 --- a/internal/core/src/common/LoadInfo.h +++ b/internal/core/src/common/LoadInfo.h @@ -21,7 +21,6 @@ #include #include "Types.h" -#include "common/CDataType.h" // NOTE: field_id can be system field // NOTE: Refer to common/SystemProperty.cpp for details @@ -38,7 +37,6 @@ struct LoadFieldDataInfo { // Set empty to disable mmap, // mmap file path will be {mmap_dir_path}/{segment_id}/{field_id} std::string mmap_dir_path = ""; - std::string url; int64_t storage_version = 0; }; diff --git a/internal/core/src/common/Span.h b/internal/core/src/common/Span.h index 950e69e003..f0d2cda4e0 100644 --- a/internal/core/src/common/Span.h +++ b/internal/core/src/common/Span.h @@ -127,7 +127,7 @@ class Span @@ -168,7 +168,7 @@ class Span< private: const embedded_type* data_; - const int64_t row_count_; - const int64_t element_sizeof_; + int64_t row_count_; + int64_t element_sizeof_; }; } // namespace milvus diff --git a/internal/core/src/common/Types.h b/internal/core/src/common/Types.h index 7dd5f462d6..17b384b3e4 100644 --- a/internal/core/src/common/Types.h +++ b/internal/core/src/common/Types.h @@ -1175,4 +1175,4 @@ struct fmt::formatter : fmt::formatter { } return fmt::formatter::format(name, ctx); } -}; \ No newline at end of file +}; diff --git a/internal/core/src/exec/expression/BinaryRangeExpr.cpp b/internal/core/src/exec/expression/BinaryRangeExpr.cpp index 07078ce579..a82938aaff 100644 --- a/internal/core/src/exec/expression/BinaryRangeExpr.cpp +++ b/internal/core/src/exec/expression/BinaryRangeExpr.cpp @@ -720,8 +720,7 @@ PhyBinaryRangeFilterExpr::ExecRangeVisitorImplForJsonForIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; if (lower_inclusive && upper_inclusive) { if (type == uint8_t(milvus::index::JSONType::STRING) || type == uint8_t(milvus::index::JSONType::DOUBLE) || diff --git a/internal/core/src/exec/expression/CompareExpr.h b/internal/core/src/exec/expression/CompareExpr.h index 5ad51a8e3b..a8e2800d54 100644 --- a/internal/core/src/exec/expression/CompareExpr.h +++ b/internal/core/src/exec/expression/CompareExpr.h @@ -295,12 +295,12 @@ class PhyCompareFilterExpr : public Expr { auto [right_chunk_id, right_chunk_offset] = get_chunk_id_and_offset(right_field_); - auto left_chunk = segment_chunk_reader_.segment_->chunk_data( + auto pw_left = segment_chunk_reader_.segment_->chunk_data( left_field_, left_chunk_id); - - auto right_chunk = - segment_chunk_reader_.segment_->chunk_data( - right_field_, right_chunk_id); + auto left_chunk = pw_left.get(); + auto pw_right = segment_chunk_reader_.segment_->chunk_data( + right_field_, right_chunk_id); + auto right_chunk = pw_right.get(); const T* left_data = left_chunk.data() + left_chunk_offset; const U* right_data = right_chunk.data() + right_chunk_offset; func.template operator()( @@ -326,10 +326,12 @@ class PhyCompareFilterExpr : public Expr { } return processed_size; } else { - auto left_chunk = + auto pw_left = segment_chunk_reader_.segment_->chunk_data(left_field_, 0); - auto right_chunk = + auto left_chunk = pw_left.get(); + auto pw_right = segment_chunk_reader_.segment_->chunk_data(right_field_, 0); + auto right_chunk = pw_right.get(); const T* left_data = left_chunk.data(); const U* right_data = right_chunk.data(); func.template operator()( @@ -363,10 +365,12 @@ class PhyCompareFilterExpr : public Expr { const auto active_count = segment_chunk_reader_.active_count_; for (size_t i = current_chunk_id_; i < num_chunk_; i++) { - auto left_chunk = + auto pw_left = segment_chunk_reader_.segment_->chunk_data(left_field_, i); - auto right_chunk = + auto left_chunk = pw_left.get(); + auto pw_right = segment_chunk_reader_.segment_->chunk_data(right_field_, i); + auto right_chunk = pw_right.get(); auto data_pos = (i == current_chunk_id_) ? current_chunk_pos_ : 0; auto size = (i == (num_chunk_ - 1)) @@ -431,10 +435,12 @@ class PhyCompareFilterExpr : public Expr { // only call this function when left and right are not indexed, so they have the same number of chunks for (size_t i = left_current_chunk_id_; i < left_num_chunk_; i++) { - auto left_chunk = + auto pw_left = segment_chunk_reader_.segment_->chunk_data(left_field_, i); - auto right_chunk = + auto left_chunk = pw_left.get(); + auto pw_right = segment_chunk_reader_.segment_->chunk_data(right_field_, i); + auto right_chunk = pw_right.get(); auto data_pos = (i == left_current_chunk_id_) ? left_current_chunk_pos_ : 0; auto size = 0; diff --git a/internal/core/src/exec/expression/Expr.cpp b/internal/core/src/exec/expression/Expr.cpp index 3f40a581eb..ba9ce61cee 100644 --- a/internal/core/src/exec/expression/Expr.cpp +++ b/internal/core/src/exec/expression/Expr.cpp @@ -33,6 +33,7 @@ #include "exec/expression/UnaryExpr.h" #include "exec/expression/ValueExpr.h" #include "expr/ITypeExpr.h" +#include "monitor/prometheus_client.h" #include diff --git a/internal/core/src/exec/expression/Expr.h b/internal/core/src/exec/expression/Expr.h index a39a8872a8..348d914f03 100644 --- a/internal/core/src/exec/expression/Expr.h +++ b/internal/core/src/exec/expression/Expr.h @@ -344,8 +344,9 @@ class SegmentExpr : public Expr { std::min(active_count_ - current_data_chunk_pos_, batch_size_); auto& skip_index = segment_->GetSkipIndex(); - auto views_info = segment_->get_batch_views( + auto pw = segment_->get_batch_views( field_id_, 0, current_data_chunk_pos_, need_size); + auto views_info = pw.get(); if (!skip_func || !skip_func(skip_index, field_id_, 0)) { // first is the raw data, second is valid_data // use valid_data to see if raw data is null @@ -381,8 +382,8 @@ class SegmentExpr : public Expr { Assert(num_data_chunk_ == 1); auto& skip_index = segment_->GetSkipIndex(); - auto [data_vec, valid_data] = - segment_->get_views_by_offsets(field_id_, 0, *input); + auto pw = segment_->get_views_by_offsets(field_id_, 0, *input); + auto [data_vec, valid_data] = pw.get(); if (!skip_func || !skip_func(skip_index, field_id_, 0)) { func(data_vec.data(), valid_data.data(), @@ -504,9 +505,9 @@ class SegmentExpr : public Expr { int64_t offset = (*input)[i]; auto [chunk_id, chunk_offset] = segment_->get_chunk_by_offset(field_id_, offset); - auto [data_vec, valid_data] = - segment_->get_views_by_offsets( - field_id_, chunk_id, {int32_t(chunk_offset)}); + auto pw = segment_->get_views_by_offsets( + field_id_, chunk_id, {int32_t(chunk_offset)}); + auto [data_vec, valid_data] = pw.get(); if (!skip_func || !skip_func(skip_index, field_id_, chunk_id)) { func.template operator()( @@ -529,7 +530,8 @@ class SegmentExpr : public Expr { int64_t offset = (*input)[i]; auto [chunk_id, chunk_offset] = segment_->get_chunk_by_offset(field_id_, offset); - auto chunk = segment_->chunk_data(field_id_, chunk_id); + auto pw = segment_->chunk_data(field_id_, chunk_id); + auto chunk = pw.get(); const T* data = chunk.data() + chunk_offset; const bool* valid_data = chunk.valid_data(); if (valid_data != nullptr) { @@ -560,7 +562,8 @@ class SegmentExpr : public Expr { return ProcessDataByOffsetsForSealedSeg( func, skip_func, input, res, valid_res, values...); } - auto chunk = segment_->chunk_data(field_id_, 0); + auto pw = segment_->chunk_data(field_id_, 0); + auto chunk = pw.get(); const T* data = chunk.data(); const bool* valid_data = chunk.valid_data(); if (!skip_func || !skip_func(skip_index, field_id_, 0)) { @@ -582,7 +585,8 @@ class SegmentExpr : public Expr { int64_t offset = (*input)[i]; auto chunk_id = offset / size_per_chunk_; auto chunk_offset = offset % size_per_chunk_; - auto chunk = segment_->chunk_data(field_id_, chunk_id); + auto pw = segment_->chunk_data(field_id_, chunk_id); + auto chunk = pw.get(); const T* data = chunk.data() + chunk_offset; const bool* valid_data = chunk.valid_data(); if (valid_data != nullptr) { @@ -642,7 +646,8 @@ class SegmentExpr : public Expr { size = std::min(size, batch_size_ - processed_size); auto& skip_index = segment_->GetSkipIndex(); - auto chunk = segment_->chunk_data(field_id_, i); + auto pw = segment_->chunk_data(field_id_, i); + auto chunk = pw.get(); const bool* valid_data = chunk.valid_data(); if (valid_data != nullptr) { valid_data += data_pos; @@ -683,14 +688,6 @@ class SegmentExpr : public Expr { ValTypes... values) { int64_t processed_size = 0; - // if constexpr (std::is_same_v || - // std::is_same_v) { - // if (segment_->type() == SegmentType::Sealed) { - // return ProcessChunkForSealedSeg( - // func, skip_func, res, values...); - // } - // } - for (size_t i = current_data_chunk_; i < num_data_chunk_; i++) { auto data_pos = (i == current_data_chunk_) ? current_data_chunk_pos_ : 0; @@ -710,9 +707,10 @@ class SegmentExpr : public Expr { if (segment_->type() == SegmentType::Sealed) { // first is the raw data, second is valid_data // use valid_data to see if raw data is null - auto [data_vec, valid_data] = - segment_->get_batch_views( - field_id_, i, data_pos, size); + auto pw = segment_->get_batch_views( + field_id_, i, data_pos, size); + auto [data_vec, valid_data] = pw.get(); + func(data_vec.data(), valid_data.data(), nullptr, @@ -724,7 +722,8 @@ class SegmentExpr : public Expr { } } if (!is_seal) { - auto chunk = segment_->chunk_data(field_id_, i); + auto pw = segment_->chunk_data(field_id_, i); + auto chunk = pw.get(); const T* data = chunk.data() + data_pos; const bool* valid_data = chunk.valid_data(); if (valid_data != nullptr) { @@ -742,15 +741,16 @@ class SegmentExpr : public Expr { const bool* valid_data; if constexpr (std::is_same_v || std::is_same_v) { - auto batch_views = segment_->get_batch_views( + auto pw = segment_->get_batch_views( field_id_, i, data_pos, size); - valid_data = batch_views.second.data(); + valid_data = pw.get().second.data(); ApplyValidData(valid_data, res + processed_size, valid_res + processed_size, size); } else { - auto chunk = segment_->chunk_data(field_id_, i); + auto pw = segment_->chunk_data(field_id_, i); + auto chunk = pw.get(); valid_data = chunk.valid_data(); if (valid_data != nullptr) { valid_data += data_pos; @@ -994,7 +994,8 @@ class SegmentExpr : public Expr { return {0, offset}; } }(); - auto chunk = segment_->chunk_data(field_id_, chunk_id); + auto pw = segment_->chunk_data(field_id_, chunk_id); + auto chunk = pw.get(); const bool* valid_data = chunk.valid_data(); if (valid_data != nullptr) { valid_result[i] = valid_data[chunk_offset]; @@ -1037,8 +1038,9 @@ class SegmentExpr : public Expr { std::is_same_v || std::is_same_v) { if (segment_->type() == SegmentType::Sealed) { - auto [data_vec, valid_data] = segment_->get_batch_views( + auto pw = segment_->get_batch_views( field_id_, i, data_pos, size); + auto [data_vec, valid_data] = pw.get(); ApplyValidData(valid_data.data(), valid_result + processed_size, valid_result + processed_size, @@ -1048,7 +1050,8 @@ class SegmentExpr : public Expr { } if (!access_sealed_variable_column) { - auto chunk = segment_->chunk_data(field_id_, i); + auto pw = segment_->chunk_data(field_id_, i); + auto chunk = pw.get(); const bool* valid_data = chunk.valid_data(); if (valid_data == nullptr) { return valid_result; diff --git a/internal/core/src/exec/expression/JsonContainsExpr.cpp b/internal/core/src/exec/expression/JsonContainsExpr.cpp index 8933a5b9a8..baec80a944 100644 --- a/internal/core/src/exec/expression/JsonContainsExpr.cpp +++ b/internal/core/src/exec/expression/JsonContainsExpr.cpp @@ -401,8 +401,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; auto array = json.array_at(offset, size); if (array.error()) { @@ -586,8 +585,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsArrayByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; auto array = json.array_at(offset, size); if (array.error()) { return false; @@ -865,8 +863,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsAllByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; auto array = json.array_at(offset, size); if (array.error()) { return false; @@ -1114,8 +1111,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsAllWithDiffTypeByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; std::set tmp_elements_index(elements_index); auto array = json.array_at(offset, size); if (array.error()) { @@ -1366,8 +1362,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsAllArrayByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; auto array = json.array_at(offset, size); if (array.error()) { return false; @@ -1602,8 +1597,7 @@ PhyJsonContainsFilterExpr::ExecJsonContainsWithDiffTypeByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; auto array = json.array_at(offset, size); if (array.error()) { return false; diff --git a/internal/core/src/exec/expression/TermExpr.cpp b/internal/core/src/exec/expression/TermExpr.cpp index 93b8fc1050..df6d188963 100644 --- a/internal/core/src/exec/expression/TermExpr.cpp +++ b/internal/core/src/exec/expression/TermExpr.cpp @@ -625,8 +625,7 @@ PhyTermFilterExpr::ExecJsonInVariableByKeyIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; if (type == uint8_t(milvus::index::JSONType::STRING) || type == uint8_t(milvus::index::JSONType::DOUBLE) || type == uint8_t(milvus::index::JSONType::INT64)) { @@ -880,7 +879,6 @@ VectorPtr PhyTermFilterExpr::ExecVisitorImplForData(EvalCtx& context) { auto* input = context.get_offset_input(); const auto& bitmap_input = context.get_bitmap_input(); - auto real_batch_size = has_offset_input_ ? input->size() : GetNextBatchSize(); if (real_batch_size == 0) { diff --git a/internal/core/src/exec/expression/UnaryExpr.cpp b/internal/core/src/exec/expression/UnaryExpr.cpp index d46d4cd941..d89a4c014d 100644 --- a/internal/core/src/exec/expression/UnaryExpr.cpp +++ b/internal/core/src/exec/expression/UnaryExpr.cpp @@ -517,91 +517,86 @@ PhyUnaryRangeFilterExpr::ExecArrayEqualForIndex(EvalCtx& context, } // cache the result to suit the framework. - auto batch_res = - ProcessIndexChunks([this, &val, reverse](Index* _) { - boost::container::vector elems; - for (auto const& element : val.array()) { - auto e = GetValueFromProto(element); - if (std::find(elems.begin(), elems.end(), e) == elems.end()) { - elems.push_back(e); - } + auto batch_res = ProcessIndexChunks([this, &val, reverse]( + Index* _) { + boost::container::vector elems; + for (auto const& element : val.array()) { + auto e = GetValueFromProto(element); + if (std::find(elems.begin(), elems.end(), e) == elems.end()) { + elems.push_back(e); } + } - // filtering by index, get candidates. - std::function retrieve; + // filtering by index, get candidates. + std::function + is_same; - // avoid use-after-free - milvus::ArrayView array_view_tmp; - if (segment_->is_chunked()) { - retrieve = [this, &array_view_tmp]( - int64_t offset) -> const milvus::ArrayView* { - auto [chunk_idx, chunk_offset] = - segment_->get_chunk_by_offset(field_id_, offset); - const auto& chunk = - segment_->template chunk_view( - field_id_, chunk_idx); - array_view_tmp = std::move(chunk.first[chunk_offset]); - return &array_view_tmp; - }; + if (segment_->is_chunked()) { + is_same = [this, reverse](milvus::proto::plan::Array& val, + int64_t offset) -> bool { + auto [chunk_idx, chunk_offset] = + segment_->get_chunk_by_offset(field_id_, offset); + auto pw = segment_->template chunk_view( + field_id_, chunk_idx); + auto chunk = pw.get(); + return chunk.first[chunk_offset].is_same_array(val) ^ reverse; + }; + } else { + auto size_per_chunk = segment_->size_per_chunk(); + is_same = [this, size_per_chunk, reverse]( + milvus::proto::plan::Array& val, + int64_t offset) -> bool { + auto chunk_idx = offset / size_per_chunk; + auto chunk_offset = offset % size_per_chunk; + auto pw = segment_->template chunk_data( + field_id_, chunk_idx); + auto chunk = pw.get(); + auto array_view = chunk.data() + chunk_offset; + return array_view->is_same_array(val) ^ reverse; + }; + } + + // collect all candidates. + std::unordered_set candidates; + std::unordered_set tmp_candidates; + auto first_callback = [&candidates](size_t offset) -> void { + candidates.insert(offset); + }; + auto callback = [&candidates, &tmp_candidates](size_t offset) -> void { + if (candidates.find(offset) != candidates.end()) { + tmp_candidates.insert(offset); + } + }; + auto execute_sub_batch = + [](Index* index_ptr, + const IndexInnerType& val, + const std::function& callback) { + index_ptr->InApplyCallback(1, &val, callback); + }; + + // run in-filter. + for (size_t idx = 0; idx < elems.size(); idx++) { + if (idx == 0) { + ProcessIndexChunksV2( + execute_sub_batch, elems[idx], first_callback); } else { - auto size_per_chunk = segment_->size_per_chunk(); - retrieve = [ size_per_chunk, this ](int64_t offset) -> auto { - auto chunk_idx = offset / size_per_chunk; - auto chunk_offset = offset % size_per_chunk; - const auto& chunk = - segment_->template chunk_data( - field_id_, chunk_idx); - return chunk.data() + chunk_offset; - }; + ProcessIndexChunksV2( + execute_sub_batch, elems[idx], callback); + candidates = std::move(tmp_candidates); } - - // compare the array via the raw data. - auto filter = [&retrieve, &val, reverse](size_t offset) -> bool { - auto data_ptr = retrieve(offset); - return data_ptr->is_same_array(val) ^ reverse; - }; - - // collect all candidates. - std::unordered_set candidates; - std::unordered_set tmp_candidates; - auto first_callback = [&candidates](size_t offset) -> void { - candidates.insert(offset); - }; - auto callback = [&candidates, - &tmp_candidates](size_t offset) -> void { - if (candidates.find(offset) != candidates.end()) { - tmp_candidates.insert(offset); - } - }; - auto execute_sub_batch = - [](Index* index_ptr, - const IndexInnerType& val, - const std::function& callback) { - index_ptr->InApplyCallback(1, &val, callback); - }; - - // run in-filter. - for (size_t idx = 0; idx < elems.size(); idx++) { - if (idx == 0) { - ProcessIndexChunksV2( - execute_sub_batch, elems[idx], first_callback); - } else { - ProcessIndexChunksV2( - execute_sub_batch, elems[idx], callback); - candidates = std::move(tmp_candidates); - } - // the size of candidates is small enough. - if (candidates.size() * 100 < active_count_) { - break; - } + // the size of candidates is small enough. + if (candidates.size() * 100 < active_count_) { + break; } - TargetBitmap res(active_count_); - // run post-filter. The filter will only be executed once in the framework. - for (const auto& candidate : candidates) { - res[candidate] = filter(candidate); - } - return res; - }); + } + TargetBitmap res(active_count_); + // run post-filter. The filter will only be executed once in the framework. + for (const auto& candidate : candidates) { + res[candidate] = is_same(val, candidate); + } + return res; + }); AssertInfo(batch_res->size() == real_batch_size, "internal error: expr processed rows {} not equal " "expect batch size {}", @@ -1224,8 +1219,7 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplJsonForIndex() { if (!json_pair.second) { return false; } - auto json = milvus::Json(json_pair.first.data(), - json_pair.first.size()); + auto& json = json_pair.first; switch (op_type) { case proto::plan::GreaterThan: if constexpr (std::is_same_v +#include "exec/expression/Utils.h" +#include "monitor/prometheus_client.h" namespace milvus { namespace exec { diff --git a/internal/core/src/exec/operator/VectorSearchNode.cpp b/internal/core/src/exec/operator/VectorSearchNode.cpp index 640e0f2da7..d4d2ee8f6a 100644 --- a/internal/core/src/exec/operator/VectorSearchNode.cpp +++ b/internal/core/src/exec/operator/VectorSearchNode.cpp @@ -16,6 +16,8 @@ #include "VectorSearchNode.h" +#include "monitor/prometheus_client.h" + namespace milvus { namespace exec { diff --git a/internal/core/src/exec/operator/groupby/SearchGroupByOperator.h b/internal/core/src/exec/operator/groupby/SearchGroupByOperator.h index b8279bbac6..5c18ac23ac 100644 --- a/internal/core/src/exec/operator/groupby/SearchGroupByOperator.h +++ b/internal/core/src/exec/operator/groupby/SearchGroupByOperator.h @@ -17,6 +17,8 @@ #pragma once #include + +#include "cachinglayer/CacheSlot.h" #include "common/QueryInfo.h" #include "common/Types.h" #include "knowhere/index/index_node.h" @@ -78,9 +80,10 @@ class SealedDataGetter : public DataGetter { const FieldId field_id_; bool from_data_; - mutable std::unordered_map> - str_view_map_; - mutable std::unordered_map> valid_map_; + mutable std::unordered_map< + int64_t, + PinWrapper, FixedVector>>> + pw_map_; // Getting str_view from segment is cpu-costly, this map is to cache this view for performance public: SealedDataGetter(const segcore::SegmentSealed& segment, FieldId& field_id) @@ -103,23 +106,22 @@ class SealedDataGetter : public DataGetter { auto chunk_id = id_offset_pair.first; auto inner_offset = id_offset_pair.second; if constexpr (std::is_same_v) { - if (str_view_map_.find(chunk_id) == str_view_map_.end()) { - auto [str_chunk_view, valid_data] = - segment_.chunk_view(field_id_, - chunk_id); - valid_map_[chunk_id] = std::move(valid_data); - str_view_map_[chunk_id] = std::move(str_chunk_view); + if (pw_map_.find(chunk_id) == pw_map_.end()) { + // for now, search_group_by does not handle null values + auto pw = segment_.chunk_view(field_id_, + chunk_id); + pw_map_[chunk_id] = std::move(pw); } - auto valid_data = valid_map_[chunk_id]; - if (!valid_data.empty()) { - if (!valid_map_[chunk_id][inner_offset]) { - return std::nullopt; - } + auto& pw = pw_map_[chunk_id]; + auto& [str_chunk_view, valid_data] = pw.get(); + if (!valid_data.empty() && !valid_data[inner_offset]) { + return std::nullopt; } - auto str_val_view = str_view_map_[chunk_id][inner_offset]; + std::string_view str_val_view = str_chunk_view[inner_offset]; return std::string(str_val_view.data(), str_val_view.length()); } else { - Span span = segment_.chunk_data(field_id_, chunk_id); + auto pw = segment_.chunk_data(field_id_, chunk_id); + auto& span = pw.get(); if (span.valid_data() && !span.valid_data()[inner_offset]) { return std::nullopt; } diff --git a/internal/core/src/futures/Executor.cpp b/internal/core/src/futures/Executor.cpp index 486ff19536..8c33123adc 100644 --- a/internal/core/src/futures/Executor.cpp +++ b/internal/core/src/futures/Executor.cpp @@ -24,7 +24,7 @@ getGlobalCPUExecutor() { static folly::CPUThreadPoolExecutor executor( thread_num, folly::CPUThreadPoolExecutor::makeDefaultPriorityQueue(kNumPriority), - std::make_shared("MILVUS_FUTURE_CPU_")); + std::make_shared("MILVUS_CPU_")); return &executor; } diff --git a/internal/core/src/index/JsonKeyStatsInvertedIndex.h b/internal/core/src/index/JsonKeyStatsInvertedIndex.h index fe358b34a1..a5f834913b 100644 --- a/internal/core/src/index/JsonKeyStatsInvertedIndex.h +++ b/internal/core/src/index/JsonKeyStatsInvertedIndex.h @@ -103,19 +103,13 @@ class JsonKeyStatsInvertedIndex : public InvertedIndexTantivy { return bitset; }; - if (is_growing) { - if (shouldTriggerCommit() || is_strong_consistency) { - if (is_data_uncommitted_) { - Commit(); - } - Reload(); - return processArray(); - } else { - return processArray(); + if (is_growing && (shouldTriggerCommit() || is_strong_consistency)) { + if (is_data_uncommitted_) { + Commit(); } - } else { - return processArray(); + Reload(); } + return processArray(); } void diff --git a/internal/core/src/index/SkipIndex.h b/internal/core/src/index/SkipIndex.h index 92ee34fd53..57a76c6fc4 100644 --- a/internal/core/src/index/SkipIndex.h +++ b/internal/core/src/index/SkipIndex.h @@ -10,13 +10,10 @@ // or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include + #include #include "common/Types.h" -#include "log/Log.h" -#include "mmap/Column.h" -#include "mmap/ChunkedColumn.h" namespace milvus { @@ -308,15 +305,15 @@ class SkipIndex { if (start > num_rows - 1) { return {std::string(), std::string(), num_rows}; } - std::string_view min_string = var_column.RawAt(start); - std::string_view max_string = var_column.RawAt(start); + std::string min_string = var_column.RawAt(start); + std::string max_string = min_string; int64_t null_count = start; for (int64_t i = start; i < num_rows; i++) { - const auto& val = var_column.RawAt(i); if (!var_column.IsValid(i)) { null_count++; continue; } + const auto& val = var_column.RawAt(i); if (val < min_string) { min_string = val; } @@ -325,7 +322,7 @@ class SkipIndex { } } // The field data may be released, so we need to copy the string to avoid invalid memory access. - return {std::string(min_string), std::string(max_string), null_count}; + return {min_string, max_string, null_count}; } private: diff --git a/internal/core/src/mmap/ChunkData.h b/internal/core/src/mmap/ChunkData.h index 42de567b0f..7f3c9f9c67 100644 --- a/internal/core/src/mmap/ChunkData.h +++ b/internal/core/src/mmap/ChunkData.h @@ -14,8 +14,10 @@ // See the License for the specific language governing permissions and // limitations under the License. #pragma once + #include "common/Array.h" #include "storage/MmapManager.h" + namespace milvus { /** * @brief FixedLengthChunk @@ -207,7 +209,7 @@ VariableLengthChunk::set( auto mcm = storage::MmapManager::GetInstance().GetMmapChunkManager(); AssertInfo( begin + length <= size_, - "failed to set a chunk with length: {} from beign {}, map_size={}", + "failed to set a chunk with length: {} from begin {}, map_size={}", length, begin, size_); diff --git a/internal/core/src/mmap/ChunkVector.h b/internal/core/src/mmap/ChunkVector.h index 3117e71e16..0fd6e2dd43 100644 --- a/internal/core/src/mmap/ChunkVector.h +++ b/internal/core/src/mmap/ChunkVector.h @@ -14,8 +14,10 @@ // See the License for the specific language governing permissions and // limitations under the License. #pragma once + #include "mmap/ChunkData.h" #include "storage/MmapManager.h" + namespace milvus { template class ChunkVectorBase { diff --git a/internal/core/src/mmap/ChunkedColumn.h b/internal/core/src/mmap/ChunkedColumn.h index f1936f7033..6fe9933ecb 100644 --- a/internal/core/src/mmap/ChunkedColumn.h +++ b/internal/core/src/mmap/ChunkedColumn.h @@ -21,82 +21,76 @@ #include #include #include -#include #include -#include #include #include #include +#include "cachinglayer/CacheSlot.h" +#include "cachinglayer/Manager.h" +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" #include "common/Array.h" #include "common/Chunk.h" -#include "common/Common.h" #include "common/EasyAssert.h" -#include "common/File.h" #include "common/FieldMeta.h" -#include "common/FieldData.h" #include "common/Span.h" -#include "fmt/format.h" -#include "log/Log.h" -#include "mmap/Utils.h" -#include "common/FieldData.h" -#include "common/FieldDataInterface.h" #include "common/Array.h" -#include "knowhere/dataset.h" -#include "monitor/prometheus_client.h" -#include "storage/MmapChunkManager.h" +#include "segcore/storagev1translator/ChunkTranslator.h" -#include "mmap/Column.h" namespace milvus { -class ChunkedColumnBase : public ColumnBase { +using namespace milvus::cachinglayer; + +std::pair inline GetChunkIDByOffset( + int64_t offset, std::vector& num_rows_until_chunk) { + AssertInfo(offset >= 0 && offset < num_rows_until_chunk.back(), + "offset is out of range, offset: {}, num rows: {}", + offset, + num_rows_until_chunk.back()); + auto iter = std::lower_bound( + num_rows_until_chunk.begin(), num_rows_until_chunk.end(), offset + 1); + size_t chunk_idx = std::distance(num_rows_until_chunk.begin(), iter) - 1; + size_t offset_in_chunk = offset - num_rows_until_chunk[chunk_idx]; + return {chunk_idx, offset_in_chunk}; +} + +class ChunkedColumnBase { public: - ChunkedColumnBase() = default; // memory mode ctor - explicit ChunkedColumnBase(const FieldMeta& field_meta) { - nullable_ = field_meta.is_nullable(); + explicit ChunkedColumnBase(std::unique_ptr> translator, + const FieldMeta& field_meta) + : nullable_(field_meta.is_nullable()), + num_chunks_(translator->num_cells()), + slot_(Manager::GetInstance().CreateCacheSlot(std::move(translator))) { + num_rows_ = GetNumRowsUntilChunk().back(); } virtual ~ChunkedColumnBase() = default; - void - AppendBatch(const FieldDataPtr data) override { - PanicInfo(ErrorCode::Unsupported, "AppendBatch not supported"); - } - - const char* - Data(int chunk_id) const override { - return chunks_[chunk_id]->Data(); - } - - virtual const char* - ValueAt(int64_t offset) const { - auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset); - return chunks_[chunk_id]->ValueAt(offset_in_chunk); - }; - - // MmappedData() returns the mmaped address - const char* - MmappedData() const override { - AssertInfo(chunks_.size() == 1, - "only support one chunk, but got {} chunk(s)", - chunks_.size()); - return chunks_[0]->RawData(); + PinWrapper + DataOfChunk(int chunk_id) const { + auto ca = SemiInlineGet(slot_->PinCells({chunk_id})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper(ca, chunk->Data()); } bool IsValid(size_t offset) const { - if (nullable_) { - auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset); - return chunks_[chunk_id]->isValid(offset_in_chunk); + if (!nullable_) { + return true; } - return true; + auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset); + return IsValid(chunk_id, offset_in_chunk); } bool IsValid(int64_t chunk_id, int64_t offset) const { if (nullable_) { - return chunks_[chunk_id]->isValid(offset); + auto ca = + SemiInlineGet(slot_->PinCells({static_cast(chunk_id)})); + auto chunk = ca->get_cell_of(chunk_id); + return chunk->isValid(offset); } return true; } @@ -113,65 +107,53 @@ class ChunkedColumnBase : public ColumnBase { int64_t num_chunks() const { - return chunks_.size(); - } - - virtual void - AddChunk(std::shared_ptr chunk) { - num_rows_until_chunk_.push_back(num_rows_); - num_rows_ += chunk->RowNums(); - chunks_.push_back(chunk); + return num_chunks_; } + // This returns only memory byte size. size_t - DataByteSize() const override { + DataByteSize() const { auto size = 0; - for (auto& chunk : chunks_) { - size += chunk->Size(); + for (auto i = 0; i < num_chunks_; i++) { + size += slot_->size_of_cell(i).memory_bytes; } return size; } int64_t chunk_row_nums(int64_t chunk_id) const { - return chunks_[chunk_id]->RowNums(); + return GetNumRowsUntilChunk(chunk_id + 1) - + GetNumRowsUntilChunk(chunk_id); } - virtual SpanBase - Span(int64_t chunk_id) const = 0; - - // used for sequential access for search - virtual BufferView - GetBatchBuffer(int64_t chunk_id, int64_t start_offset, int64_t length) { + virtual PinWrapper + Span(int64_t chunk_id) const { PanicInfo(ErrorCode::Unsupported, - "GetBatchBuffer only supported for VariableColumn"); + "Span only supported for ChunkedColumn"); } - virtual std::string_view - RawAt(const size_t i) const { - PanicInfo(ErrorCode::Unsupported, - "RawAt only supported for VariableColumn"); - } - - virtual std::pair, FixedVector> + virtual PinWrapper< + std::pair, FixedVector>> StringViews(int64_t chunk_id, - std::optional> offset_len) const { + std::optional> offset_len = + std::nullopt) const { PanicInfo(ErrorCode::Unsupported, "StringViews only supported for VariableColumn"); } - virtual std::pair, FixedVector> + virtual PinWrapper, FixedVector>> ArrayViews(int64_t chunk_id, std::optional> offset_len) const { PanicInfo(ErrorCode::Unsupported, "ArrayViews only supported for ArrayChunkedColumn"); } - virtual std::pair, FixedVector> + virtual PinWrapper< + std::pair, FixedVector>> ViewsByOffsets(int64_t chunk_id, const FixedVector& offsets) const { PanicInfo(ErrorCode::Unsupported, - "viewsbyoffsets only supported for VariableColumn"); + "ViewsByOffsets only supported for VariableColumn"); } std::pair @@ -180,268 +162,148 @@ class ChunkedColumnBase : public ColumnBase { "offset {} is out of range, num_rows: {}", offset, num_rows_); - - auto iter = std::lower_bound(num_rows_until_chunk_.begin(), - num_rows_until_chunk_.end(), - offset + 1); - size_t chunk_idx = - std::distance(num_rows_until_chunk_.begin(), iter) - 1; - size_t offset_in_chunk = offset - num_rows_until_chunk_[chunk_idx]; - return {chunk_idx, offset_in_chunk}; + auto num_rows_until_chunk = GetNumRowsUntilChunk(); + return ::milvus::GetChunkIDByOffset(offset, num_rows_until_chunk); } - std::shared_ptr + PinWrapper GetChunk(int64_t chunk_id) const { - return chunks_[chunk_id]; + auto ca = SemiInlineGet(slot_->PinCells({chunk_id})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper(ca, chunk); } int64_t GetNumRowsUntilChunk(int64_t chunk_id) const { - return num_rows_until_chunk_[chunk_id]; + return GetNumRowsUntilChunk()[chunk_id]; } const std::vector& GetNumRowsUntilChunk() const { - return num_rows_until_chunk_; + auto meta = static_cast( + slot_->meta()); + return meta->num_rows_until_chunk_; } protected: bool nullable_{false}; size_t num_rows_{0}; - std::vector num_rows_until_chunk_; - - private: - // void - // UpdateMetricWhenMmap(size_t mmaped_size) { - // UpdateMetricWhenMmap(mapping_type_, mmaped_size); - // } - - // void - // UpdateMetricWhenMmap(bool is_map_anonymous, size_t mapped_size) { - // if (mapping_type_ == MappingType::MAP_WITH_ANONYMOUS) { - // milvus::monitor::internal_mmap_allocated_space_bytes_anon.Observe( - // mapped_size); - // milvus::monitor::internal_mmap_in_used_space_bytes_anon.Increment( - // mapped_size); - // } else { - // milvus::monitor::internal_mmap_allocated_space_bytes_file.Observe( - // mapped_size); - // milvus::monitor::internal_mmap_in_used_space_bytes_file.Increment( - // mapped_size); - // } - // } - - // void - // UpdateMetricWhenMunmap(size_t mapped_size) { - // if (mapping_type_ == MappingType::MAP_WITH_ANONYMOUS) { - // milvus::monitor::internal_mmap_in_used_space_bytes_anon.Decrement( - // mapped_size); - // } else { - // milvus::monitor::internal_mmap_in_used_space_bytes_file.Decrement( - // mapped_size); - // } - // } - - private: - storage::MmapChunkManagerPtr mcm_ = nullptr; - - protected: - std::vector> chunks_; + size_t num_chunks_{0}; + mutable std::shared_ptr> slot_; }; class ChunkedColumn : public ChunkedColumnBase { public: - ChunkedColumn() = default; // memory mode ctor - explicit ChunkedColumn(const FieldMeta& field_meta) - : ChunkedColumnBase(field_meta) { + explicit ChunkedColumn(std::unique_ptr> translator, + const FieldMeta& field_meta) + : ChunkedColumnBase(std::move(translator), field_meta) { } - explicit ChunkedColumn(const FieldMeta& field_meta, - const std::vector>& chunks) - : ChunkedColumnBase(field_meta) { - for (auto& chunk : chunks) { - AddChunk(chunk); - } + // TODO(tiered storage 1): this method should be replaced with a bulk access method. + const char* + ValueAt(int64_t offset) { + auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset); + auto ca = + SemiInlineGet(slot_->PinCells({static_cast(chunk_id)})); + auto chunk = ca->get_cell_of(chunk_id); + return chunk->ValueAt(offset_in_chunk); } - ~ChunkedColumn() override = default; - - SpanBase + PinWrapper Span(int64_t chunk_id) const override { - return std::static_pointer_cast(chunks_[chunk_id]) - ->Span(); + auto ca = SemiInlineGet(slot_->PinCells({chunk_id})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper( + ca, static_cast(chunk)->Span()); } }; -// when mmap is used, size_, data_ and num_rows_ of ColumnBase are used. -class ChunkedSparseFloatColumn : public ChunkedColumnBase { - public: - // memory mode ctor - explicit ChunkedSparseFloatColumn(const FieldMeta& field_meta) - : ChunkedColumnBase(field_meta) { - } - - explicit ChunkedSparseFloatColumn( - const FieldMeta& field_meta, - const std::vector>& chunks) - : ChunkedColumnBase(field_meta) { - for (auto& chunk : chunks) { - AddChunk(chunk); - } - } - - ~ChunkedSparseFloatColumn() override = default; - - void - AddChunk(std::shared_ptr chunk) override { - num_rows_until_chunk_.push_back(num_rows_); - num_rows_ += chunk->RowNums(); - chunks_.push_back(chunk); - dim_ = std::max( - dim_, - std::static_pointer_cast(chunk)->Dim()); - } - - SpanBase - Span(int64_t chunk_id) const override { - PanicInfo(ErrorCode::Unsupported, - "Span not supported for sparse float column"); - } - - int64_t - Dim() const { - return dim_; - } - - private: - int64_t dim_ = 0; -}; - template class ChunkedVariableColumn : public ChunkedColumnBase { public: - using ViewType = - std::conditional_t, std::string_view, T>; + static_assert( + std::is_same_v || std::is_same_v, + "ChunkedVariableColumn only supports std::string or Json types"); // memory mode ctor - explicit ChunkedVariableColumn(const FieldMeta& field_meta) - : ChunkedColumnBase(field_meta) { - } - explicit ChunkedVariableColumn( - const FieldMeta& field_meta, - const std::vector>& chunks) - : ChunkedColumnBase(field_meta) { - for (auto& chunk : chunks) { - AddChunk(chunk); - } + std::unique_ptr> translator, + const FieldMeta& field_meta) + : ChunkedColumnBase(std::move(translator), field_meta) { } - ~ChunkedVariableColumn() override = default; - - SpanBase - Span(int64_t chunk_id) const override { - PanicInfo(ErrorCode::NotImplemented, - "span() interface is not implemented for variable column"); - } - - std::pair, FixedVector> + PinWrapper, FixedVector>> StringViews(int64_t chunk_id, std::optional> offset_len = std::nullopt) const override { - return std::static_pointer_cast(chunks_[chunk_id]) - ->StringViews(offset_len); + auto ca = SemiInlineGet(slot_->PinCells({chunk_id})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper< + std::pair, FixedVector>>( + ca, static_cast(chunk)->StringViews(offset_len)); } - std::pair, FixedVector> + PinWrapper, FixedVector>> ViewsByOffsets(int64_t chunk_id, const FixedVector& offsets) const override { - return std::static_pointer_cast(chunks_[chunk_id]) - ->ViewsByOffsets(offsets); + auto ca = SemiInlineGet(slot_->PinCells({chunk_id})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper< + std::pair, FixedVector>>( + ca, static_cast(chunk)->ViewsByOffsets(offsets)); } - BufferView - GetBatchBuffer(int64_t chunk_id, - int64_t start_offset, - int64_t length) override { - BufferView buffer_view; - std::vector elements; - elements.push_back( - {chunks_[chunk_id]->Data(), - std::static_pointer_cast(chunks_[chunk_id]) - ->Offsets(), - static_cast(start_offset), - static_cast(start_offset + length)}); - - buffer_view.data_ = elements; - return buffer_view; - } - - ViewType - operator[](const int i) const { + // TODO(tiered storage 1): this method should be replaced with a bulk access method. + // RawAt is called in three cases: + // 1. bulk_subscript, pass in an offset array, access the specified rows. + // 2. load, create skip index or text index, access all rows. (SkipIndex.h and CreateTextIndex) + // 3. GetJsonData, json related index will use this, see if it can be modified to batch access or batch pin. + T + RawAt(const size_t i) const { if (i < 0 || i > num_rows_) { PanicInfo(ErrorCode::OutOfRange, "index out of range"); } auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i); + auto ca = + SemiInlineGet(slot_->PinCells({static_cast(chunk_id)})); + auto chunk = ca->get_cell_of(chunk_id); std::string_view str_view = - std::static_pointer_cast(chunks_[chunk_id]) - -> - operator[](offset_in_chunk); - return ViewType(str_view.data(), str_view.size()); - } - - std::string_view - RawAt(const size_t i) const { - return std::string_view((*this)[i]); + static_cast(chunk)->operator[](offset_in_chunk); + return T(str_view.data(), str_view.size()); } }; + class ChunkedArrayColumn : public ChunkedColumnBase { public: // memory mode ctor - explicit ChunkedArrayColumn(const FieldMeta& field_meta) - : ChunkedColumnBase(field_meta) { - } - - explicit ChunkedArrayColumn( - const FieldMeta& field_meta, - const std::vector>& chunks) - : ChunkedColumnBase(field_meta) { - for (auto& chunk : chunks) { - AddChunk(chunk); - } - } - - ~ChunkedArrayColumn() override = default; - - SpanBase - Span(int64_t chunk_id) const override { - PanicInfo(ErrorCode::NotImplemented, - "span() interface is not implemented for arr chunk column"); - } - - ArrayView - operator[](const int i) const { - auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i); - return std::static_pointer_cast(chunks_[chunk_id]) - ->View(offset_in_chunk); + explicit ChunkedArrayColumn(std::unique_ptr> translator, + const FieldMeta& field_meta) + : ChunkedColumnBase(std::move(translator), field_meta) { } + // TODO(tiered storage 1): this method should be replaced with a bulk access method. ScalarArray RawAt(const int i) const { auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i); - return std::static_pointer_cast(chunks_[chunk_id]) + auto ca = + SemiInlineGet(slot_->PinCells({static_cast(chunk_id)})); + auto chunk = ca->get_cell_of(chunk_id); + return static_cast(chunk) ->View(offset_in_chunk) .output_data(); } - std::pair, FixedVector> + PinWrapper, FixedVector>> ArrayViews(int64_t chunk_id, std::optional> offset_len = std::nullopt) const override { - return std::static_pointer_cast(chunks_[chunk_id]) - ->Views(offset_len); + auto ca = + SemiInlineGet(slot_->PinCells({static_cast(chunk_id)})); + auto chunk = ca->get_cell_of(chunk_id); + return PinWrapper, FixedVector>>( + ca, static_cast(chunk)->Views(offset_len)); } }; -} // namespace milvus \ No newline at end of file +} // namespace milvus diff --git a/internal/core/src/mmap/Column.h b/internal/core/src/mmap/Column.h deleted file mode 100644 index dd875f1422..0000000000 --- a/internal/core/src/mmap/Column.h +++ /dev/null @@ -1,137 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "common/Array.h" -#include "common/Common.h" -#include "common/EasyAssert.h" -#include "common/File.h" -#include "common/FieldMeta.h" -#include "common/FieldData.h" -#include "common/Span.h" -#include "fmt/format.h" -#include "log/Log.h" -#include "mmap/Utils.h" -#include "common/FieldData.h" -#include "common/FieldDataInterface.h" -#include "common/Array.h" -#include "knowhere/dataset.h" -#include "monitor/prometheus_client.h" -#include "storage/MmapChunkManager.h" - -namespace milvus { - -constexpr size_t DEFAULT_PK_VRCOL_BLOCK_SIZE = 1; -constexpr size_t DEFAULT_MEM_VRCOL_BLOCK_SIZE = 32; -constexpr size_t DEFAULT_MMAP_VRCOL_BLOCK_SIZE = 256; - -/** - * ColumnBase and its subclasses are designed to store and retrieve the raw data - * of a field. - * - * It has 3 types of constructors corresponding to 3 MappingTypes: - * - * 1. MAP_WITH_ANONYMOUS: ColumnBase(size_t reserve_size, const FieldMeta& field_meta) - * - * This is used when we store the entire data in memory. Upon return, a piece - * of unwritten memory is allocated and the caller can fill the memory with data by - * calling AppendBatch/Append. - * - * 2. MAP_WITH_FILE: ColumnBase(const File& file, size_t size, const FieldMeta& field_meta) - * - * This is used when the raw data has already been written into a file, and we - * simply mmap the file to memory and interpret the memory as a column. In this - * mode, since the data is already in the file/mmapped memory, calling AppendBatch - * and Append is not allowed. - * - * 3. MAP_WITH_MANAGER: ColumnBase(size_t reserve, - * const DataType& data_type, - * storage::MmapChunkManagerPtr mcm, - * storage::MmapChunkDescriptorPtr descriptor, - * bool nullable) - * - * This is used when we want to mmap but don't want to download all the data at once. - * Instead, we download the data in chunks, cache and mmap each chunk as a single - * ColumnBase. Upon return, a piece of unwritten mmaped memory is allocated by the chunk - * manager, and the caller should fill the memory with data by calling AppendBatch - * and Append. - * - * - Types with fixed length can use the Column subclass directly. - * - Types with variable lengths: - * - SparseFloatColumn: - * - To store sparse float vectors. - * - All 3 modes are supported. - * - VariableColumn: - * - To store string like types such as VARCHAR and JSON. - * - MAP_WITH_MANAGER is not yet supported(as of 2024.09.11). - * - ArrayColumn: - * - To store ARRAY types. - * - MAP_WITH_MANAGER is not yet supported(as of 2024.09.11). - * - */ -class ColumnBase { - /** - * - data_ points at a piece of memory of size data_cap_size_ + padding_. - * Based on mapping_type_, such memory can be: - * - an anonymous memory region, allocated by mmap(MAP_ANON) - * - a file-backed memory region, mapped by mmap(MAP_FILE) - * - a memory region managed by MmapChunkManager, allocated by - * MmapChunkManager::Allocate() - * - * Memory Layout of `data_`: - * - * |<-- data_cap_size_ -->|<-- padding_ -->| - * |<-- data_size_ -->|<-- free space -->| - * - * AppendBatch/Append should first check if there's enough space for new data. - * If not, call ExpandData() to expand the space. - * - * - only the first data_cap_size_ bytes can be used to store actual data. - * - padding at the end is to ensure when all values are empty, we don't try - * to allocate/mmap 0 bytes memory, which will cause mmap() to fail. - * - data_size_ is the number of bytes currently used to store actual data. - * - num_rows_ is the number of rows currently stored. - * - valid_data_ is a FixedVector indicating whether each element is - * not null. it is only used when nullable is true. - * - nullable_ is true if null(0 byte) is a valid value for the column. - * - */ - public: - virtual size_t - DataByteSize() const = 0; - - virtual const char* - MmappedData() const = 0; - - virtual void - AppendBatch(const FieldDataPtr data) = 0; - - virtual const char* - Data(int chunk_id) const = 0; -}; -} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/mmap/Types.h b/internal/core/src/mmap/Types.h index 7613425635..50ecff08cc 100644 --- a/internal/core/src/mmap/Types.h +++ b/internal/core/src/mmap/Types.h @@ -19,9 +19,8 @@ #include #include #include -#include "arrow/record_batch.h" + #include "common/FieldData.h" -#include "storage/DataCodec.h" namespace milvus { @@ -39,18 +38,6 @@ struct FieldDataInfo { arrow_reader_channel = std::make_shared(); } - FieldDataInfo( - int64_t field_id, - size_t row_count, - const std::vector>& batch) - : field_id(field_id), row_count(row_count) { - arrow_reader_channel = std::make_shared(); - for (auto& data : batch) { - arrow_reader_channel->push(data); - } - arrow_reader_channel->close(); - } - int64_t field_id; size_t row_count; std::string mmap_dir_path; diff --git a/internal/core/src/monitor/prometheus_client.cpp b/internal/core/src/monitor/prometheus_client.cpp index 7ec3dadae2..de34cc0a2d 100644 --- a/internal/core/src/monitor/prometheus_client.cpp +++ b/internal/core/src/monitor/prometheus_client.cpp @@ -326,4 +326,218 @@ DEFINE_PROMETHEUS_GAUGE(internal_cgo_executing_task_total_all, internal_cgo_executing_task_total, {}); +// --- caching layer metrics --- +// TODO(tiered storage 1): choose better buckets. + +std::map cacheMemoryLabel = {{"location", "memory"}}; +std::map cacheDiskLabel = {{"location", "disk"}}; +std::map cacheMixedLabel = {{"location", "mixed"}}; + +// Cache slot count +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_slot_count, + "[cpp]cache slot count"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_slot_count_memory, + internal_cache_slot_count, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_slot_count_disk, + internal_cache_slot_count, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_slot_count_mixed, + internal_cache_slot_count, + cacheMixedLabel); + +// Cache cell count +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_cell_count, + "[cpp]cache cell count"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_count_memory, + internal_cache_cell_count, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_count_disk, + internal_cache_cell_count, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_count_mixed, + internal_cache_cell_count, + cacheMixedLabel); + +// Cache cell loaded count +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_cell_loaded_count, + "[cpp]cache cell loaded count"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_memory, + internal_cache_cell_loaded_count, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_disk, + internal_cache_cell_loaded_count, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_mixed, + internal_cache_cell_loaded_count, + cacheMixedLabel); + +// Cache load latency histogram +DEFINE_PROMETHEUS_HISTOGRAM_FAMILY(internal_cache_load_latency, + "[cpp]cache load latency histogram"); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS(internal_cache_load_latency_memory, + internal_cache_load_latency, + cacheMemoryLabel, + secondsBuckets); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS(internal_cache_load_latency_disk, + internal_cache_load_latency, + cacheDiskLabel, + secondsBuckets); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS(internal_cache_load_latency_mixed, + internal_cache_load_latency, + cacheMixedLabel, + secondsBuckets); + +// Cache hit rate (represented by hit/miss counters) +std::map cacheHitMemoryLabels = { + {"result", "hit"}, {"location", "memory"}}; +std::map cacheHitDiskLabels = {{"result", "hit"}, + {"location", "disk"}}; +std::map cacheHitMixedLabels = { + {"result", "hit"}, {"location", "mixed"}}; +std::map cacheMissMemoryLabels = { + {"result", "miss"}, {"location", "memory"}}; +std::map cacheMissDiskLabels = {{"result", "miss"}, + {"location", "disk"}}; +std::map cacheMissMixedLabels = { + {"result", "miss"}, {"location", "mixed"}}; +DEFINE_PROMETHEUS_COUNTER_FAMILY(internal_cache_op_result_count, + "[cpp]cache operation result count"); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_memory, + internal_cache_op_result_count, + cacheHitMemoryLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_disk, + internal_cache_op_result_count, + cacheHitDiskLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_mixed, + internal_cache_op_result_count, + cacheHitMixedLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_memory, + internal_cache_op_result_count, + cacheMissMemoryLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_disk, + internal_cache_op_result_count, + cacheMissDiskLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_mixed, + internal_cache_op_result_count, + cacheMissMixedLabels); + +// Cache usage (bytes) +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_used_bytes, + "[cpp]currently used bytes in cache"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_used_bytes_memory, + internal_cache_used_bytes, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_used_bytes_disk, + internal_cache_used_bytes, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_used_bytes_mixed, + internal_cache_used_bytes, + cacheMixedLabel); + +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_capacity_bytes, + "[cpp]total capacity bytes of cache"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_capacity_bytes_memory, + internal_cache_capacity_bytes, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_capacity_bytes_disk, + internal_cache_capacity_bytes, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_capacity_bytes_mixed, + internal_cache_capacity_bytes, + cacheMixedLabel); + +// Eviction count and resource size +DEFINE_PROMETHEUS_COUNTER_FAMILY(internal_cache_eviction_count, + "[cpp]cache eviction count"); +DEFINE_PROMETHEUS_COUNTER(internal_cache_eviction_count_memory, + internal_cache_eviction_count, + cacheMemoryLabel); +DEFINE_PROMETHEUS_COUNTER(internal_cache_eviction_count_disk, + internal_cache_eviction_count, + cacheDiskLabel); +DEFINE_PROMETHEUS_COUNTER(internal_cache_eviction_count_mixed, + internal_cache_eviction_count, + cacheMixedLabel); + +DEFINE_PROMETHEUS_COUNTER_FAMILY(internal_cache_evicted_bytes, + "[cpp]total bytes evicted from cache"); +DEFINE_PROMETHEUS_COUNTER(internal_cache_evicted_bytes_memory, + internal_cache_evicted_bytes, + cacheMemoryLabel); +DEFINE_PROMETHEUS_COUNTER(internal_cache_evicted_bytes_disk, + internal_cache_evicted_bytes, + cacheDiskLabel); +DEFINE_PROMETHEUS_COUNTER(internal_cache_evicted_bytes_mixed, + internal_cache_evicted_bytes, + cacheMixedLabel); + +// Cache item lifetime histogram +DEFINE_PROMETHEUS_HISTOGRAM_FAMILY(internal_cache_item_lifetime_seconds, + "[cpp]cache item lifetime histogram"); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS( + internal_cache_item_lifetime_seconds_memory, + internal_cache_item_lifetime_seconds, + cacheMemoryLabel, + secondsBuckets); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS( + internal_cache_item_lifetime_seconds_disk, + internal_cache_item_lifetime_seconds, + cacheDiskLabel, + secondsBuckets); +DEFINE_PROMETHEUS_HISTOGRAM_WITH_BUCKETS( + internal_cache_item_lifetime_seconds_mixed, + internal_cache_item_lifetime_seconds, + cacheMixedLabel, + secondsBuckets); + +// Load error rate (represented by success/fail counters) +std::map cacheLoadSuccessMemoryLabels = { + {"status", "success"}, {"location", "memory"}}; +std::map cacheLoadSuccessDiskLabels = { + {"status", "success"}, {"location", "disk"}}; +std::map cacheLoadSuccessMixedLabels = { + {"status", "success"}, {"location", "mixed"}}; +std::map cacheLoadFailMemoryLabels = { + {"status", "fail"}, {"location", "memory"}}; +std::map cacheLoadFailDiskLabels = { + {"status", "fail"}, {"location", "disk"}}; +std::map cacheLoadFailMixedLabels = { + {"status", "fail"}, {"location", "mixed"}}; +DEFINE_PROMETHEUS_COUNTER_FAMILY(internal_cache_load_count, + "[cpp]cache load operation count"); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_success_memory, + internal_cache_load_count, + cacheLoadSuccessMemoryLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_success_disk, + internal_cache_load_count, + cacheLoadSuccessDiskLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_success_mixed, + internal_cache_load_count, + cacheLoadSuccessMixedLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_memory, + internal_cache_load_count, + cacheLoadFailMemoryLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_disk, + internal_cache_load_count, + cacheLoadFailDiskLabels); +DEFINE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_mixed, + internal_cache_load_count, + cacheLoadFailMixedLabels); + +// Cache system memory overhead (bytes) +DEFINE_PROMETHEUS_GAUGE_FAMILY(internal_cache_memory_overhead_bytes, + "[cpp]cache system memory overhead in bytes"); +DEFINE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_memory, + internal_cache_memory_overhead_bytes, + cacheMemoryLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_disk, + internal_cache_memory_overhead_bytes, + cacheDiskLabel); +DEFINE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_mixed, + internal_cache_memory_overhead_bytes, + cacheMixedLabel); + +// --- caching layer metrics end --- + } // namespace milvus::monitor diff --git a/internal/core/src/monitor/prometheus_client.h b/internal/core/src/monitor/prometheus_client.h index 579f9f562e..ad33129ee9 100644 --- a/internal/core/src/monitor/prometheus_client.h +++ b/internal/core/src/monitor/prometheus_client.h @@ -158,4 +158,72 @@ DECLARE_PROMETHEUS_GAUGE(internal_cgo_inflight_task_total_all); DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cgo_executing_task_total); DECLARE_PROMETHEUS_GAUGE(internal_cgo_executing_task_total_all); +// --- caching layer metrics --- + +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_slot_count); +DECLARE_PROMETHEUS_GAUGE(internal_cache_slot_count_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_slot_count_disk); +DECLARE_PROMETHEUS_GAUGE(internal_cache_slot_count_mixed); + +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_cell_count); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_count_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_count_disk); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_count_mixed); + +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_cell_loaded_count); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_disk); +DECLARE_PROMETHEUS_GAUGE(internal_cache_cell_loaded_count_mixed); + +DECLARE_PROMETHEUS_HISTOGRAM_FAMILY(internal_cache_load_latency); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_load_latency_memory); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_load_latency_disk); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_load_latency_mixed); + +DECLARE_PROMETHEUS_COUNTER_FAMILY(internal_cache_op_result_count); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_disk); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_hit_mixed); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_disk); +DECLARE_PROMETHEUS_COUNTER(internal_cache_op_result_count_miss_mixed); + +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_used_bytes); +DECLARE_PROMETHEUS_GAUGE(internal_cache_used_bytes_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_used_bytes_disk); + +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_capacity_bytes); +DECLARE_PROMETHEUS_GAUGE(internal_cache_capacity_bytes_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_capacity_bytes_disk); + +DECLARE_PROMETHEUS_COUNTER_FAMILY(internal_cache_eviction_count); +DECLARE_PROMETHEUS_COUNTER(internal_cache_eviction_count_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_eviction_count_disk); +DECLARE_PROMETHEUS_COUNTER(internal_cache_eviction_count_mixed); + +DECLARE_PROMETHEUS_COUNTER_FAMILY(internal_cache_evicted_bytes); +DECLARE_PROMETHEUS_COUNTER(internal_cache_evicted_bytes_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_evicted_bytes_disk); + +DECLARE_PROMETHEUS_HISTOGRAM_FAMILY(internal_cache_item_lifetime_seconds); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_item_lifetime_seconds_memory); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_item_lifetime_seconds_disk); +DECLARE_PROMETHEUS_HISTOGRAM(internal_cache_item_lifetime_seconds_mixed); + +DECLARE_PROMETHEUS_COUNTER_FAMILY(internal_cache_load_count); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_success_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_success_disk); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_success_mixed); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_memory); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_disk); +DECLARE_PROMETHEUS_COUNTER(internal_cache_load_count_fail_mixed); + +// TODO(tiered storage 1): not added +DECLARE_PROMETHEUS_GAUGE_FAMILY(internal_cache_memory_overhead_bytes); +DECLARE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_memory); +DECLARE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_disk); +DECLARE_PROMETHEUS_GAUGE(internal_cache_memory_overhead_bytes_mixed); + +// --- caching layer metrics end --- + } // namespace milvus::monitor diff --git a/internal/core/src/query/CachedSearchIterator.cpp b/internal/core/src/query/CachedSearchIterator.cpp index e39571442e..d99a6d39df 100644 --- a/internal/core/src/query/CachedSearchIterator.cpp +++ b/internal/core/src/query/CachedSearchIterator.cpp @@ -9,9 +9,11 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License +#include + +#include "mmap/ChunkedColumn.h" #include "query/CachedSearchIterator.h" #include "query/SearchBruteForce.h" -#include namespace milvus::query { @@ -72,9 +74,10 @@ CachedSearchIterator::InitializeChunkedIterators( int64_t offset = 0; chunked_heaps_.resize(nq_); for (int64_t chunk_id = 0; chunk_id < num_chunks_; ++chunk_id) { + // TODO(tiered storage 1): this should store PinWrapper. Double check all places. auto [chunk_data, chunk_size] = get_chunk_data(chunk_id); auto sub_data = query::dataset::RawDataset{ - offset, query_ds.dim, chunk_size, chunk_data}; + offset, query_ds.dim, chunk_size, chunk_data.get()}; auto expected_iterators = GetBruteForceSearchIterators( query_ds, sub_data, search_info, index_info, bitset, data_type); @@ -121,17 +124,17 @@ CachedSearchIterator::CachedSearchIterator( index_info, bitset, data_type, - [&vec_data, vec_size_per_chunk, row_count]( - int64_t chunk_id) -> std::pair { - const auto chunk_data = vec_data->get_chunk_data(chunk_id); + [&vec_data, vec_size_per_chunk, row_count](int64_t chunk_id) { + const void* chunk_data = vec_data->get_chunk_data(chunk_id); + auto pw = milvus::cachinglayer::PinWrapper(chunk_data); int64_t chunk_size = std::min( vec_size_per_chunk, row_count - chunk_id * vec_size_per_chunk); - return {chunk_data, chunk_size}; + return std::make_pair(pw, chunk_size); }); } CachedSearchIterator::CachedSearchIterator( - const std::shared_ptr& column, + ChunkedColumnBase* column, const dataset::SearchDataset& query_ds, const SearchInfo& search_info, const std::map& index_info, @@ -153,11 +156,11 @@ CachedSearchIterator::CachedSearchIterator( index_info, bitset, data_type, - [&column](int64_t chunk_id) { - const char* chunk_data = column->Data(chunk_id); + [column](int64_t chunk_id) { + auto pw = column->DataOfChunk(chunk_id).transform( + [](const auto& x) { return static_cast(x); }); int64_t chunk_size = column->chunk_row_nums(chunk_id); - return std::make_pair(static_cast(chunk_data), - chunk_size); + return std::make_pair(pw, chunk_size); }); } diff --git a/internal/core/src/query/CachedSearchIterator.h b/internal/core/src/query/CachedSearchIterator.h index 1ede6d0c12..5378f943b9 100644 --- a/internal/core/src/query/CachedSearchIterator.h +++ b/internal/core/src/query/CachedSearchIterator.h @@ -12,12 +12,15 @@ #pragma once #include + +#include "cachinglayer/CacheSlot.h" #include "common/BitsetView.h" #include "common/QueryInfo.h" #include "common/QueryResult.h" #include "query/helper.h" #include "segcore/ConcurrentVector.h" #include "index/VectorIndex.h" +#include "mmap/ChunkedColumn.h" namespace milvus::query { @@ -58,7 +61,7 @@ class CachedSearchIterator { const milvus::DataType& data_type); // For sealed segment with chunked data, BF - CachedSearchIterator(const std::shared_ptr& column, + CachedSearchIterator(ChunkedColumnBase* column, const dataset::SearchDataset& dataset, const SearchInfo& search_info, const std::map& index_info, @@ -83,7 +86,8 @@ class CachedSearchIterator { using IterIdx = size_t; using IterIdDisIdPair = std::pair; using GetChunkDataFunc = - std::function(int64_t)>; + std::function, + int64_t>(int64_t)>; int64_t batch_size_ = 0; std::vector iterators_; diff --git a/internal/core/src/query/SearchOnSealed.cpp b/internal/core/src/query/SearchOnSealed.cpp index 7dddd3d40c..e2f37487b8 100644 --- a/internal/core/src/query/SearchOnSealed.cpp +++ b/internal/core/src/query/SearchOnSealed.cpp @@ -17,7 +17,6 @@ #include "common/BitsetView.h" #include "common/QueryInfo.h" #include "common/Types.h" -#include "mmap/Column.h" #include "query/CachedSearchIterator.h" #include "query/SearchBruteForce.h" #include "query/SearchOnSealed.h" @@ -86,15 +85,15 @@ SearchOnSealedIndex(const Schema& schema, } void -SearchOnSealed(const Schema& schema, - std::shared_ptr column, - const SearchInfo& search_info, - const std::map& index_info, - const void* query_data, - int64_t num_queries, - int64_t row_count, - const BitsetView& bitview, - SearchResult& result) { +SearchOnSealedColumn(const Schema& schema, + ChunkedColumnBase* column, + const SearchInfo& search_info, + const std::map& index_info, + const void* query_data, + int64_t num_queries, + int64_t row_count, + const BitsetView& bitview, + SearchResult& result) { auto field_id = search_info.field_id_; auto& field = schema[field_id]; @@ -129,7 +128,8 @@ SearchOnSealed(const Schema& schema, auto offset = 0; for (int i = 0; i < num_chunk; ++i) { - auto vec_data = column->Data(i); + auto pw = column->DataOfChunk(i); + auto vec_data = pw.get(); auto chunk_size = column->chunk_row_nums(i); auto raw_dataset = query::dataset::RawDataset{offset, dim, chunk_size, vec_data}; @@ -167,15 +167,15 @@ SearchOnSealed(const Schema& schema, } void -SearchOnSealed(const Schema& schema, - const void* vec_data, - const SearchInfo& search_info, - const std::map& index_info, - const void* query_data, - int64_t num_queries, - int64_t row_count, - const BitsetView& bitset, - SearchResult& result) { +SearchOnSealedData(const Schema& schema, + const void* vec_data, + const SearchInfo& search_info, + const std::map& index_info, + const void* query_data, + int64_t num_queries, + int64_t row_count, + const BitsetView& bitset, + SearchResult& result) { auto field_id = search_info.field_id_; auto& field = schema[field_id]; diff --git a/internal/core/src/query/SearchOnSealed.h b/internal/core/src/query/SearchOnSealed.h index b3254b1c14..547747624c 100644 --- a/internal/core/src/query/SearchOnSealed.h +++ b/internal/core/src/query/SearchOnSealed.h @@ -28,25 +28,25 @@ SearchOnSealedIndex(const Schema& schema, SearchResult& search_result); void -SearchOnSealed(const Schema& schema, - std::shared_ptr column, - const SearchInfo& search_info, - const std::map& index_info, - const void* query_data, - int64_t num_queries, - int64_t row_count, - const BitsetView& bitset, - SearchResult& result); +SearchOnSealedColumn(const Schema& schema, + ChunkedColumnBase* column, + const SearchInfo& search_info, + const std::map& index_info, + const void* query_data, + int64_t num_queries, + int64_t row_count, + const BitsetView& bitset, + SearchResult& result); void -SearchOnSealed(const Schema& schema, - const void* vec_data, - const SearchInfo& search_info, - const std::map& index_info, - const void* query_data, - int64_t num_queries, - int64_t row_count, - const BitsetView& bitset, - SearchResult& result); +SearchOnSealedData(const Schema& schema, + const void* vec_data, + const SearchInfo& search_info, + const std::map& index_info, + const void* query_data, + int64_t num_queries, + int64_t row_count, + const BitsetView& bitset, + SearchResult& result); } // namespace milvus::query diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp index c83ad2d462..0a589dc7cd 100644 --- a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp +++ b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp @@ -18,7 +18,6 @@ #include #include -#include #include #include #include @@ -28,14 +27,12 @@ #include "Utils.h" #include "Types.h" +#include "cachinglayer/Manager.h" #include "common/Array.h" #include "common/Chunk.h" -#include "common/ChunkWriter.h" #include "common/Consts.h" #include "common/EasyAssert.h" -#include "common/FieldData.h" #include "common/FieldMeta.h" -#include "common/File.h" #include "common/Json.h" #include "common/LoadInfo.h" #include "common/Schema.h" @@ -44,20 +41,21 @@ #include "google/protobuf/message_lite.h" #include "index/VectorMemIndex.h" #include "mmap/ChunkedColumn.h" -#include "mmap/Utils.h" #include "mmap/Types.h" +#include "monitor/prometheus_client.h" #include "log/Log.h" #include "pb/schema.pb.h" -#include "query/ScalarIndex.h" -#include "query/SearchBruteForce.h" #include "query/SearchOnSealed.h" -#include "storage/DataCodec.h" +#include "segcore/storagev1translator/ChunkTranslator.h" +#include "segcore/storagev1translator/DefaultValueChunkTranslator.h" +#include "segcore/storagev1translator/InsertRecordTranslator.h" #include "storage/Util.h" -#include "storage/ThreadPools.h" #include "storage/MmapManager.h" namespace milvus::segcore { +using namespace milvus::cachinglayer; + static inline void set_bit(BitsetType& bitset, FieldId field_id, bool flag = true) { auto pos = field_id.get() - START_USER_FIELDID; @@ -121,7 +119,6 @@ ChunkedSegmentSealedImpl::LoadVecIndex(const LoadIndexInfo& info) { set_bit(binlog_index_bitset_, field_id, false); vector_indexings_.drop_field_indexing(field_id); } - update_row_count(row_count); vector_indexings_.append_field_indexing( field_id, metric_type, @@ -132,45 +129,16 @@ ChunkedSegmentSealedImpl::LoadVecIndex(const LoadIndexInfo& info) { id_); } -void -ChunkedSegmentSealedImpl::WarmupChunkCache(const FieldId field_id, - bool mmap_enabled) { - auto& field_meta = schema_->operator[](field_id); - AssertInfo(field_meta.is_vector(), "vector field is not vector type"); - - if (!get_bit(index_ready_bitset_, field_id) && - !get_bit(binlog_index_bitset_, field_id)) { - return; - } - - AssertInfo(vector_indexings_.is_ready(field_id), - "vector index is not ready"); - auto field_indexing = vector_indexings_.get_field_indexing(field_id); - auto vec_index = - dynamic_cast(field_indexing->indexing_.get()); - AssertInfo(vec_index, "invalid vector indexing"); - - auto it = field_data_info_.field_infos.find(field_id.get()); - AssertInfo(it != field_data_info_.field_infos.end(), - "cannot find binlog file for field: {}, seg: {}", - field_id.get(), - id_); - auto field_info = it->second; - - auto cc = storage::MmapManager::GetInstance().GetChunkCache(); - for (const auto& data_path : field_info.insert_files) { - auto column = cc->Read(data_path, field_meta, mmap_enabled, true); - } -} - void ChunkedSegmentSealedImpl::LoadScalarIndex(const LoadIndexInfo& info) { // NOTE: lock only when data is ready to avoid starvation auto field_id = FieldId(info.field_id); auto& field_meta = schema_->operator[](field_id); + auto is_pk = field_id == schema_->get_primary_field_id(); + // if segment is pk sorted, user created indexes bring no performance gain but extra memory usage - if (is_sorted_by_pk_ && field_id == schema_->get_primary_field_id()) { + if (is_pk && is_sorted_by_pk_) { LOG_INFO( "segment pk sorted, skip user index loading for primary key field"); return; @@ -205,24 +173,23 @@ ChunkedSegmentSealedImpl::LoadScalarIndex(const LoadIndexInfo& info) { std::move(const_cast(info).index); set_bit(index_ready_bitset_, field_id, true); - update_row_count(row_count); // release field column if the index contains raw data // only release non-primary field when in pk sorted mode if (scalar_indexings_[field_id]->HasRawData() && - get_bit(field_data_ready_bitset_, field_id) && - (schema_->get_primary_field_id() != field_id || !is_sorted_by_pk_)) { + get_bit(field_data_ready_bitset_, field_id) && !is_pk) { + // We do not erase the primary key field: if insert record is evicted from memory, when reloading it'll + // need the pk field again. fields_.erase(field_id); set_bit(field_data_ready_bitset_, field_id, false); } - - lck.unlock(); } void ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) { - // NOTE: lock only when data is ready to avoid starvation - // only one field for now, parallel load field data in golang size_t num_rows = storage::GetNumRowsForLoadInfo(load_info); + AssertInfo( + !num_rows_.has_value() || num_rows_ == num_rows, + "num_rows_ is set but not equal to num_rows of LoadFieldDataInfo"); for (auto& [id, info] : load_info.field_infos) { AssertInfo(info.row_count > 0, "The row count of field data is 0"); @@ -243,381 +210,121 @@ ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) { field_id.get(), num_rows); - auto parallel_degree = static_cast( - DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE); - field_data_info.arrow_reader_channel->set_capacity(parallel_degree * 2); - auto& pool = - ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE); - pool.Submit(LoadArrowReaderFromRemote, + if (SystemProperty::Instance().IsSystem(field_id)) { + if (field_id == RowFieldID) { + // ignore row id field + continue; + } + std::unique_ptr>> translator = + std::make_unique( + this->get_segment_id(), + DataType::INT64, + field_data_info, + schema_, + is_sorted_by_pk_, insert_files, - field_data_info.arrow_reader_channel); - - LOG_INFO("segment {} submits load field {} task to thread pool", - this->get_segment_id(), - field_id.get()); - bool use_mmap = false; - if (!info.enable_mmap || - SystemProperty::Instance().IsSystem(field_id)) { - LoadFieldData(field_id, field_data_info); - } else { - MapFieldData(field_id, field_data_info); - use_mmap = true; - } - LOG_INFO("segment {} loads field {} mmap {} done", - this->get_segment_id(), - field_id.get(), - use_mmap); - } -} - -void -ChunkedSegmentSealedImpl::LoadFieldData(FieldId field_id, FieldDataInfo& data) { - auto num_rows = data.row_count; - if (SystemProperty::Instance().IsSystem(field_id)) { - auto system_field_type = - SystemProperty::Instance().GetSystemFieldType(field_id); - if (system_field_type == SystemFieldType::Timestamp) { - std::vector timestamps(num_rows); - int64_t offset = 0; - FieldMeta field_meta(FieldName(""), - FieldId(0), - DataType::INT64, - false, - std::nullopt); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = std::dynamic_pointer_cast( - create_chunk(field_meta, 1, array_vec)); - std::copy_n(static_cast(chunk->Span().data()), - chunk->Span().row_count(), - timestamps.data() + offset); - offset += chunk->Span().row_count(); + this); + insert_record_slot_ = + Manager::GetInstance().CreateCacheSlot(std::move(translator)); + deleted_record_ = std::make_unique>( + insert_record_slot_, + [this](const PkType& pk, Timestamp timestamp) { + return this->search_pk(pk, timestamp); + }, + this->get_segment_id()); + { + std::unique_lock lck(mutex_); + update_row_count(num_rows); } - - // for (auto& data : field_data) { - // int64_t row_count = data->get_num_rows(); - // std::copy_n(static_cast(data->Data()), - // row_count, - // timestamps.data() + offset); - // offset += row_count; - // } - - TimestampIndex index; - auto min_slice_length = num_rows < 4096 ? 1 : 4096; - auto meta = GenerateFakeSlices( - timestamps.data(), num_rows, min_slice_length); - index.set_length_meta(std::move(meta)); - // todo ::opt to avoid copy timestamps from field data - index.build_with(timestamps.data(), num_rows); - - // use special index - std::unique_lock lck(mutex_); - AssertInfo(insert_record_.timestamps_.empty(), "already exists"); - insert_record_.timestamps_.set_data_raw( - 0, timestamps.data(), timestamps.size()); - insert_record_.timestamp_index_ = std::move(index); - AssertInfo(insert_record_.timestamps_.num_chunk() == 1, - "num chunk not equal to 1 for sealed segment"); - stats_.mem_size += sizeof(Timestamp) * data.row_count; + ++system_ready_count_; } else { - AssertInfo(system_field_type == SystemFieldType::RowId, - "System field type of id column is not RowId"); - // Consume rowid field data but not really load it - // storage::CollectFieldDataChannel(data.arrow_reader_channel); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - } - } - ++system_ready_count_; - } else { - // prepare data - auto& field_meta = (*schema_)[field_id]; - auto data_type = field_meta.get_data_type(); + auto field_meta = schema_->operator[](field_id); + std::unique_ptr> translator = + std::make_unique( + this->get_segment_id(), + field_meta, + field_data_info, + insert_files, + info.enable_mmap); - // Don't allow raw data and index exist at the same time - // AssertInfo(!get_bit(index_ready_bitset_, field_id), - // "field data can't be loaded when indexing exists"); + std::shared_ptr column{}; - std::shared_ptr column{}; - if (IsVariableDataType(data_type)) { - int64_t field_data_size = 0; + auto data_type = field_meta.get_data_type(); switch (data_type) { case milvus::DataType::STRING: case milvus::DataType::VARCHAR: case milvus::DataType::TEXT: { - auto var_column = + column = std::make_shared>( - field_meta); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = create_chunk(field_meta, 1, array_vec); - var_column->AddChunk(chunk); - } - // var_column->Seal(); - field_data_size = var_column->DataByteSize(); - stats_.mem_size += var_column->DataByteSize(); - LoadStringSkipIndex(field_id, 0, *var_column); - column = std::move(var_column); + std::move(translator), field_meta); break; } case milvus::DataType::JSON: { - auto var_column = + column = std::make_shared>( - field_meta); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = create_chunk(field_meta, 1, array_vec); - var_column->AddChunk(chunk); - } - // var_column->Seal(); - stats_.mem_size += var_column->DataByteSize(); - field_data_size = var_column->DataByteSize(); - column = std::move(var_column); + std::move(translator), field_meta); break; } case milvus::DataType::ARRAY: { - auto var_column = - std::make_shared(field_meta); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = create_chunk(field_meta, 1, array_vec); - var_column->AddChunk(chunk); - } - column = std::move(var_column); - break; - } - case milvus::DataType::VECTOR_SPARSE_FLOAT: { - auto col = - std::make_shared(field_meta); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = create_chunk(field_meta, 1, array_vec); - col->AddChunk(chunk); - } - column = std::move(col); + column = std::make_shared( + std::move(translator), field_meta); break; } default: { - PanicInfo(DataTypeInvalid, - fmt::format("unsupported data type", data_type)); + column = std::make_shared( + std::move(translator), field_meta); + break; } } - // update average row data size - SegmentInternalInterface::set_field_avg_size( - field_id, num_rows, field_data_size); - } else { - column = std::make_shared(field_meta); - std::shared_ptr r; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = - read_single_column_batches(r->reader); - auto chunk = - create_chunk(field_meta, - IsVectorDataType(field_meta.get_data_type()) && - !IsSparseFloatVectorDataType( - field_meta.get_data_type()) - ? field_meta.get_dim() - : 1, - array_vec); - // column->AppendBatch(field_data); - // stats_.mem_size += field_data->Size(); - column->AddChunk(chunk); + { + std::unique_lock lck(mutex_); + fields_.emplace(field_id, column); + if (info.enable_mmap) { + mmap_fields_.insert(field_id); + } + if (!num_rows_.has_value()) { + num_rows_ = num_rows; + } } - auto num_chunk = column->num_chunks(); - for (int i = 0; i < num_chunk; ++i) { - LoadPrimitiveSkipIndex(field_id, - i, - data_type, - column->Span(i).data(), - column->Span(i).valid_data(), - column->Span(i).row_count()); + if (!info.enable_mmap) { + stats_.mem_size += column->DataByteSize(); + if (IsVariableDataType(data_type)) { + if (data_type == milvus::DataType::STRING || + data_type == milvus::DataType::VARCHAR || + data_type == milvus::DataType::TEXT) { + auto var_column = std::dynamic_pointer_cast< + ChunkedVariableColumn>(column); + LoadStringSkipIndex(field_id, 0, *var_column); + } + // update average row data size + SegmentInternalInterface::set_field_avg_size( + field_id, num_rows, column->DataByteSize()); + } else { + auto num_chunk = column->num_chunks(); + for (int i = 0; i < num_chunk; ++i) { + auto pw = column->Span(i); + LoadPrimitiveSkipIndex(field_id, + i, + data_type, + pw.get().data(), + pw.get().valid_data(), + pw.get().row_count()); + } + } + } + + if (generate_interim_index(field_id)) { + std::unique_lock lck(mutex_); + // mmap_fields is useless, no change + fields_.erase(field_id); + set_bit(field_data_ready_bitset_, field_id, false); + } else { + std::unique_lock lck(mutex_); + set_bit(field_data_ready_bitset_, field_id, true); } } - - AssertInfo(column->NumRows() == num_rows, - fmt::format("data lost while loading column {}: loaded " - "num rows {} but expected {}", - data.field_id, - column->NumRows(), - num_rows)); - - { - std::unique_lock lck(mutex_); - fields_.emplace(field_id, column); - } - - // set pks to offset - if (schema_->get_primary_field_id() == field_id && !is_sorted_by_pk_) { - AssertInfo(field_id.get() != -1, "Primary key is -1"); - AssertInfo(insert_record_.empty_pks(), "already exists"); - insert_record_.insert_pks(data_type, column); - insert_record_.seal_pks(); - } - - bool use_temp_index = false; - { - // update num_rows to build temperate binlog index - std::unique_lock lck(mutex_); - update_row_count(num_rows); - } - - if (generate_interim_index(field_id)) { - std::unique_lock lck(mutex_); - fields_.erase(field_id); - set_bit(field_data_ready_bitset_, field_id, false); - use_temp_index = true; - } - - if (!use_temp_index) { - std::unique_lock lck(mutex_); - set_bit(field_data_ready_bitset_, field_id, true); - } - } - { - std::unique_lock lck(mutex_); - update_row_count(num_rows); - } -} - -void -ChunkedSegmentSealedImpl::MapFieldData(const FieldId field_id, - FieldDataInfo& data) { - auto filepath = std::filesystem::path(data.mmap_dir_path) / - std::to_string(get_segment_id()) / - std::to_string(field_id.get()); - auto dir = filepath.parent_path(); - std::filesystem::create_directories(dir); - - auto file = File::Open(filepath.string(), O_CREAT | O_TRUNC | O_RDWR); - - auto& field_meta = (*schema_)[field_id]; - auto data_type = field_meta.get_data_type(); - - // write the field data to disk - std::vector indices{}; - std::vector> element_indices{}; - // FixedVector valid_data{}; - std::shared_ptr r; - - size_t file_offset = 0; - std::vector> chunks; - while (data.arrow_reader_channel->pop(r)) { - arrow::ArrayVector array_vec = read_single_column_batches(r->reader); - auto chunk = create_chunk( - field_meta, - IsVectorDataType(field_meta.get_data_type()) && - !IsSparseFloatVectorDataType(field_meta.get_data_type()) - ? field_meta.get_dim() - : 1, - file, - file_offset, - array_vec); - file_offset += chunk->Size(); - chunks.push_back(chunk); - } - // WriteFieldPadding(file, data_type, total_written); - std::shared_ptr column{}; - auto num_rows = data.row_count; - if (IsVariableDataType(data_type)) { - switch (data_type) { - case milvus::DataType::STRING: - case milvus::DataType::VARCHAR: - case milvus::DataType::TEXT: { - // auto var_column = std::make_shared>( - // file, - // total_written, - // field_meta, - // DEFAULT_MMAP_VRCOL_BLOCK_SIZE); - auto var_column = - std::make_shared>( - field_meta, chunks); - // var_column->Seal(std::move(indices)); - column = std::move(var_column); - break; - } - case milvus::DataType::JSON: { - auto var_column = - std::make_shared>( - field_meta, chunks); - // var_column->Seal(std::move(indices)); - column = std::move(var_column); - break; - } - case milvus::DataType::ARRAY: { - auto arr_column = - std::make_shared(field_meta, chunks); - // arr_column->Seal(std::move(indices), - // std::move(element_indices)); - column = std::move(arr_column); - break; - } - case milvus::DataType::VECTOR_SPARSE_FLOAT: { - auto sparse_column = std::make_shared( - field_meta, chunks); - // sparse_column->Seal(std::move(indices)); - column = std::move(sparse_column); - break; - } - default: { - PanicInfo(DataTypeInvalid, - fmt::format("unsupported data type {}", data_type)); - } - } - } else { - column = std::make_shared(field_meta, chunks); - } - - // column->SetValidData(std::move(valid_data)); - - { - std::unique_lock lck(mutex_); - fields_.emplace(field_id, column); - mmap_fields_.insert(field_id); - } - - { - std::unique_lock lck(mutex_); - update_row_count(num_rows); - } - - auto ok = unlink(filepath.c_str()); - AssertInfo(ok == 0, - fmt::format("failed to unlink mmap data file {}, err: {}", - filepath.c_str(), - strerror(errno))); - - // set pks to offset - if (schema_->get_primary_field_id() == field_id && !is_sorted_by_pk_) { - AssertInfo(field_id.get() != -1, "Primary key is -1"); - AssertInfo(insert_record_.empty_pks(), "already exists"); - insert_record_.insert_pks(data_type, column); - insert_record_.seal_pks(); - } - - bool use_interim_index = false; - if (generate_interim_index(field_id)) { - std::unique_lock lck(mutex_); - // mmap_fields is useless, no change - fields_.erase(field_id); - set_bit(field_data_ready_bitset_, field_id, false); - use_interim_index = true; - } - - if (!use_interim_index) { - std::unique_lock lck(mutex_); - set_bit(field_data_ready_bitset_, field_id, true); } } @@ -636,7 +343,7 @@ ChunkedSegmentSealedImpl::LoadDeletedRecord(const LoadDeletedRecordInfo& info) { auto timestamps = reinterpret_cast(info.timestamps); // step 2: push delete info to delete_record - deleted_record_.LoadPush(pks, timestamps); + deleted_record_->LoadPush(pks, timestamps); } void @@ -701,54 +408,26 @@ ChunkedSegmentSealedImpl::num_rows_until_chunk(FieldId field_id, return fields_.at(field_id)->GetNumRowsUntilChunk(chunk_id); } -std::pair> -ChunkedSegmentSealedImpl::get_chunk_buffer(FieldId field_id, - int64_t chunk_id, - int64_t start_offset, - int64_t length) const { - std::shared_lock lck(mutex_); - AssertInfo(get_bit(field_data_ready_bitset_, field_id), - "Can't get bitset element at " + std::to_string(field_id.get())); - if (auto it = fields_.find(field_id); it != fields_.end()) { - auto& field_data = it->second; - FixedVector valid_data; - if (field_data->IsNullable()) { - valid_data.reserve(length); - for (int i = 0; i < length; i++) { - valid_data.push_back( - field_data->IsValid(chunk_id, start_offset + i)); - } - } - return std::make_pair( - field_data->GetBatchBuffer(chunk_id, start_offset, length), - valid_data); - } - PanicInfo(ErrorCode::UnexpectedError, - "get_chunk_buffer only used for variable column field"); -} - bool ChunkedSegmentSealedImpl::is_mmap_field(FieldId field_id) const { std::shared_lock lck(mutex_); return mmap_fields_.find(field_id) != mmap_fields_.end(); } -SpanBase +PinWrapper ChunkedSegmentSealedImpl::chunk_data_impl(FieldId field_id, int64_t chunk_id) const { std::shared_lock lck(mutex_); AssertInfo(get_bit(field_data_ready_bitset_, field_id), "Can't get bitset element at " + std::to_string(field_id.get())); if (auto it = fields_.find(field_id); it != fields_.end()) { - auto& field_data = it->second; - return field_data->Span(chunk_id); + return it->second->Span(chunk_id); } - auto field_data = insert_record_.get_data_base(field_id); - // system field - return field_data->get_span_base(0); + PanicInfo(ErrorCode::UnexpectedError, + "chunk_data_impl only used for chunk column field "); } -std::pair, FixedVector> +PinWrapper, FixedVector>> ChunkedSegmentSealedImpl::chunk_array_view_impl( FieldId field_id, int64_t chunk_id, @@ -758,14 +437,13 @@ ChunkedSegmentSealedImpl::chunk_array_view_impl( AssertInfo(get_bit(field_data_ready_bitset_, field_id), "Can't get bitset element at " + std::to_string(field_id.get())); if (auto it = fields_.find(field_id); it != fields_.end()) { - auto& field_data = it->second; - return field_data->ArrayViews(chunk_id, offset_len); + return it->second->ArrayViews(chunk_id, offset_len); } PanicInfo(ErrorCode::UnexpectedError, "chunk_array_view_impl only used for chunk column field "); } -std::pair, FixedVector> +PinWrapper, FixedVector>> ChunkedSegmentSealedImpl::chunk_string_view_impl( FieldId field_id, int64_t chunk_id, @@ -775,14 +453,14 @@ ChunkedSegmentSealedImpl::chunk_string_view_impl( AssertInfo(get_bit(field_data_ready_bitset_, field_id), "Can't get bitset element at " + std::to_string(field_id.get())); if (auto it = fields_.find(field_id); it != fields_.end()) { - auto& field_data = it->second; - return field_data->StringViews(chunk_id, offset_len); + auto column = it->second; + return column->StringViews(chunk_id, offset_len); } PanicInfo(ErrorCode::UnexpectedError, "chunk_string_view_impl only used for variable column field "); } -std::pair, FixedVector> +PinWrapper, FixedVector>> ChunkedSegmentSealedImpl::chunk_view_by_offsets( FieldId field_id, int64_t chunk_id, @@ -791,8 +469,7 @@ ChunkedSegmentSealedImpl::chunk_view_by_offsets( AssertInfo(get_bit(field_data_ready_bitset_, field_id), "Can't get bitset element at " + std::to_string(field_id.get())); if (auto it = fields_.find(field_id); it != fields_.end()) { - auto& field_data = it->second; - return field_data->ViewsByOffsets(chunk_id, offsets); + return it->second->ViewsByOffsets(chunk_id, offsets); } PanicInfo(ErrorCode::UnexpectedError, "chunk_view_by_offsets only used for variable column field "); @@ -804,8 +481,7 @@ ChunkedSegmentSealedImpl::chunk_index_impl(FieldId field_id, AssertInfo(scalar_indexings_.find(field_id) != scalar_indexings_.end(), "Cannot find scalar_indexing with field_id: " + std::to_string(field_id.get())); - auto ptr = scalar_indexings_.at(field_id).get(); - return ptr; + return scalar_indexings_.at(field_id).get(); } int64_t @@ -817,7 +493,7 @@ ChunkedSegmentSealedImpl::get_row_count() const { int64_t ChunkedSegmentSealedImpl::get_deleted_count() const { std::shared_lock lck(mutex_); - return deleted_record_.size(); + return deleted_record_->size(); } const Schema& @@ -829,7 +505,7 @@ void ChunkedSegmentSealedImpl::mask_with_delete(BitsetTypeView& bitset, int64_t ins_barrier, Timestamp timestamp) const { - deleted_record_.Query(bitset, ins_barrier, timestamp); + deleted_record_->Query(bitset, ins_barrier, timestamp); } void @@ -891,59 +567,19 @@ ChunkedSegmentSealedImpl::vector_search(SearchInfo& search_info, col_index_meta_->GetFieldIndexMeta(field_id).GetIndexParams(); } - query::SearchOnSealed(*schema_, - vec_data, - search_info, - index_info, - query_data, - query_count, - row_count, - bitset, - output); + query::SearchOnSealedColumn(*schema_, + vec_data.get(), + search_info, + index_info, + query_data, + query_count, + row_count, + bitset, + output); milvus::tracer::AddEvent("finish_searching_vector_data"); } } -std::tuple -ChunkedSegmentSealedImpl::GetFieldDataPath(FieldId field_id, - int64_t offset) const { - auto offset_in_binlog = offset; - auto data_path = std::string(); - auto it = field_data_info_.field_infos.find(field_id.get()); - AssertInfo(it != field_data_info_.field_infos.end(), - fmt::format("cannot find binlog file for field: {}, seg: {}", - field_id.get(), - id_)); - auto field_info = it->second; - - for (auto i = 0; i < field_info.insert_files.size(); i++) { - if (offset_in_binlog < field_info.entries_nums[i]) { - data_path = field_info.insert_files[i]; - break; - } else { - offset_in_binlog -= field_info.entries_nums[i]; - } - } - return {data_path, offset_in_binlog}; -} - -std::tuple< - std::string, - std::shared_ptr< - ChunkedColumnBase>> static ReadFromChunkCache(const storage:: - ChunkCachePtr& cc, - const std::string& - data_path, - const storage:: - MmapChunkDescriptorPtr& - descriptor, - const FieldMeta& - field_meta) { - auto column = cc->Read(data_path, field_meta, true); - cc->Prefetch(data_path); - return {data_path, std::dynamic_pointer_cast(column)}; -} - std::unique_ptr ChunkedSegmentSealedImpl::get_vector(FieldId field_id, const int64_t* ids, @@ -982,106 +618,23 @@ ChunkedSegmentSealedImpl::get_vector(FieldId field_id, } } - // If index doesn't have raw data, get vector from chunk cache. - auto cc = storage::MmapManager::GetInstance().GetChunkCache(); - - // group by data_path - auto id_to_data_path = - std::unordered_map>{}; - auto path_to_column = - std::unordered_map>{}; - for (auto i = 0; i < count; i++) { - const auto& tuple = GetFieldDataPath(field_id, ids[i]); - id_to_data_path.emplace(ids[i], tuple); - path_to_column.emplace(std::get<0>(tuple), nullptr); - } - - // read and prefetch - auto& pool = ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::HIGH); - std::vector>>> - futures; - futures.reserve(path_to_column.size()); - for (const auto& iter : path_to_column) { - const auto& data_path = iter.first; - futures.emplace_back(pool.Submit( - ReadFromChunkCache, cc, data_path, mmap_descriptor_, field_meta)); - } - - for (auto& future : futures) { - const auto& [data_path, column] = future.get(); - path_to_column[data_path] = column; - } - - if (field_meta.get_data_type() == DataType::VECTOR_SPARSE_FLOAT) { - auto buf = std::vector>(count); - for (auto i = 0; i < count; ++i) { - const auto& [data_path, offset_in_binlog] = - id_to_data_path.at(ids[i]); - const auto& column = path_to_column.at(data_path); - AssertInfo( - offset_in_binlog < column->NumRows(), - "column idx out of range, idx: {}, size: {}, data_path: {}", - offset_in_binlog, - column->NumRows(), - data_path); - auto sparse_column = - std::dynamic_pointer_cast(column); - AssertInfo(sparse_column, "incorrect column created"); - buf[i] = *static_cast*>( - static_cast( - sparse_column->ValueAt(offset_in_binlog))); - } - return segcore::CreateVectorDataArrayFrom( - buf.data(), count, field_meta); - } else { - // assign to data array - auto row_bytes = field_meta.get_sizeof(); - auto buf = std::vector(count * row_bytes); - for (auto i = 0; i < count; ++i) { - AssertInfo(id_to_data_path.count(ids[i]) != 0, "id not found"); - const auto& [data_path, offset_in_binlog] = - id_to_data_path.at(ids[i]); - AssertInfo(path_to_column.count(data_path) != 0, - "column not found"); - const auto& column = path_to_column.at(data_path); - AssertInfo( - offset_in_binlog * row_bytes < column->DataByteSize(), - "column idx out of range, idx: {}, size: {}, data_path: {}", - offset_in_binlog * row_bytes, - column->DataByteSize(), - data_path); - auto vector = column->ValueAt(offset_in_binlog); - std::memcpy(buf.data() + i * row_bytes, vector, row_bytes); - } - return segcore::CreateVectorDataArrayFrom( - buf.data(), count, field_meta); - } + AssertInfo(false, "get_vector called on vector index without raw data"); + return nullptr; } void ChunkedSegmentSealedImpl::DropFieldData(const FieldId field_id) { - if (SystemProperty::Instance().IsSystem(field_id)) { - auto system_field_type = - SystemProperty::Instance().GetSystemFieldType(field_id); - - std::unique_lock lck(mutex_); - --system_ready_count_; - if (system_field_type == SystemFieldType::Timestamp) { - insert_record_.timestamps_.clear(); - } - lck.unlock(); - } else { - std::unique_lock lck(mutex_); - if (get_bit(field_data_ready_bitset_, field_id)) { - fields_.erase(field_id); - set_bit(field_data_ready_bitset_, field_id, false); - } - if (get_bit(binlog_index_bitset_, field_id)) { - set_bit(binlog_index_bitset_, field_id, false); - vector_indexings_.drop_field_indexing(field_id); - } - lck.unlock(); + AssertInfo(!SystemProperty::Instance().IsSystem(field_id), + "Dropping system field is not supported, field id: {}", + field_id.get()); + std::unique_lock lck(mutex_); + if (get_bit(field_data_ready_bitset_, field_id)) { + fields_.erase(field_id); + set_bit(field_data_ready_bitset_, field_id, false); + } + if (get_bit(binlog_index_bitset_, field_id)) { + set_bit(binlog_index_bitset_, field_id, false); + vector_indexings_.drop_field_indexing(field_id); } } @@ -1107,11 +660,11 @@ ChunkedSegmentSealedImpl::check_search(const query::Plan* plan) const { "Extra info of search plan doesn't have value"); if (!is_system_field_ready()) { - PanicInfo( - FieldNotLoaded, - "failed to load row ID or timestamp, potential missing bin logs or " - "empty segments. Segment ID = " + - std::to_string(this->id_)); + PanicInfo(FieldNotLoaded, + "failed to load row ID or timestamp, potential missing " + "bin logs or " + "empty segments. Segment ID = " + + std::to_string(this->id_)); } auto& request_fields = plan->extra_info_opt_.value().involved_fields_; @@ -1142,22 +695,13 @@ ChunkedSegmentSealedImpl::check_search(const query::Plan* plan) const { std::vector ChunkedSegmentSealedImpl::search_pk(const PkType& pk, Timestamp timestamp) const { + auto ir_accessor = pin_insert_record(); + auto ir = ir_accessor->get_cell_of(0); if (!is_sorted_by_pk_) { - return insert_record_.search_pk(pk, timestamp); + return ir->search_pk(pk, timestamp); } - return search_sorted_pk(pk, [this, timestamp](int64_t offset) { - return insert_record_.timestamps_[offset] <= timestamp; - }); -} - -std::vector -ChunkedSegmentSealedImpl::search_pk(const PkType& pk, - int64_t insert_barrier) const { - if (!is_sorted_by_pk_) { - return insert_record_.search_pk(pk, insert_barrier); - } - return search_sorted_pk(pk, [insert_barrier](int64_t offset) { - return offset < insert_barrier; + return search_sorted_pk(pk, [this, timestamp, ir](int64_t offset) { + return ir->timestamps()[offset] <= timestamp; }); } @@ -1177,7 +721,8 @@ ChunkedSegmentSealedImpl::search_sorted_pk(const PkType& pk, auto num_chunk = pk_column->num_chunks(); for (int i = 0; i < num_chunk; ++i) { - auto src = reinterpret_cast(pk_column->Data(i)); + auto pw = pk_column->DataOfChunk(i); + auto src = reinterpret_cast(pw.get()); auto chunk_row_num = pk_column->chunk_row_nums(i); auto it = std::lower_bound( src, @@ -1200,16 +745,13 @@ ChunkedSegmentSealedImpl::search_sorted_pk(const PkType& pk, case DataType::VARCHAR: { auto target = std::get(pk); // get varchar pks - auto var_column = - std::dynamic_pointer_cast>( - pk_column); - auto num_chunk = var_column->num_chunks(); + auto num_chunk = pk_column->num_chunks(); for (int i = 0; i < num_chunk; ++i) { // TODO @xiaocai2333, @sunby: chunk need to record the min/max. auto num_rows_until_chunk = pk_column->GetNumRowsUntilChunk(i); - auto string_chunk = std::dynamic_pointer_cast( - var_column->GetChunk(i)); + auto pw = pk_column->GetChunk(i); + auto string_chunk = static_cast(pw.get()); auto offset = string_chunk->binary_search_string(target); for (; offset != -1 && offset < string_chunk->RowNums() && string_chunk->operator[](offset) == target; @@ -1238,7 +780,9 @@ std::pair, bool> ChunkedSegmentSealedImpl::find_first(int64_t limit, const BitsetType& bitset) const { if (!is_sorted_by_pk_) { - return insert_record_.pk2offset_->find_first(limit, bitset); + auto ir_accessor = pin_insert_record(); + return ir_accessor->get_cell_of(0)->pk2offset_->find_first(limit, + bitset); } if (limit == Unlimited || limit == NoLimit) { limit = num_rows_.value(); @@ -1280,18 +824,13 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( index_ready_bitset_(schema->size()), binlog_index_bitset_(schema->size()), scalar_indexings_(schema->size()), - insert_record_(*schema, MAX_ROW_COUNT), + insert_record_slot_(nullptr), schema_(schema), id_(segment_id), col_index_meta_(index_meta), TEST_skip_index_for_retrieve_(TEST_skip_index_for_retrieve), is_sorted_by_pk_(is_sorted_by_pk), - deleted_record_( - &insert_record_, - [this](const PkType& pk, Timestamp timestamp) { - return this->search_pk(pk, timestamp); - }, - segment_id) { + deleted_record_(nullptr) { mmap_descriptor_ = std::shared_ptr( new storage::MmapChunkDescriptor({segment_id, SegmentType::Sealed})); auto mcm = storage::MmapManager::GetInstance().GetMmapChunkManager(); @@ -1299,16 +838,6 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( } ChunkedSegmentSealedImpl::~ChunkedSegmentSealedImpl() { - auto cc = storage::MmapManager::GetInstance().GetChunkCache(); - if (cc == nullptr) { - return; - } - // munmap and remove binlog from chunk cache - for (const auto& iter : field_data_info_.field_infos) { - for (const auto& binlog : iter.second.insert_files) { - cc->Remove(binlog); - } - } if (mmap_descriptor_ != nullptr) { auto mm = storage::MmapManager::GetInstance().GetMmapChunkManager(); mm->UnRegister(mmap_descriptor_); @@ -1323,13 +852,14 @@ ChunkedSegmentSealedImpl::bulk_subscript(SystemFieldType system_type, AssertInfo(is_system_field_ready(), "System field isn't ready when do bulk_insert, segID:{}", id_); + auto ir_accessor = pin_insert_record(); switch (system_type) { case SystemFieldType::Timestamp: AssertInfo( - insert_record_.timestamps_.num_chunk() == 1, + ir_accessor->get_cell_of(0)->timestamps().num_chunk() == 1, "num chunk of timestamp not equal to 1 for sealed segment"); bulk_subscript_impl( - this->insert_record_.timestamps_.get_chunk_data(0), + ir_accessor->get_cell_of(0)->timestamps().get_chunk_data(0), seg_offsets, count, static_cast(output)); @@ -1358,40 +888,27 @@ ChunkedSegmentSealedImpl::bulk_subscript_impl(const void* src_raw, } template void -ChunkedSegmentSealedImpl::bulk_subscript_impl(const ChunkedColumnBase* field, +ChunkedSegmentSealedImpl::bulk_subscript_impl(ChunkedColumnBase* field, const int64_t* seg_offsets, int64_t count, T* dst) { static_assert(IsScalar); + auto column = reinterpret_cast(field); for (int64_t i = 0; i < count; ++i) { auto offset = seg_offsets[i]; dst[i] = *static_cast( - static_cast(field->ValueAt(offset))); - } -} - -template -void -ChunkedSegmentSealedImpl::bulk_subscript_impl(const ChunkedColumnBase* column, - const int64_t* seg_offsets, - int64_t count, - void* dst_raw) { - auto field = reinterpret_cast*>(column); - auto dst = reinterpret_cast(dst_raw); - for (int64_t i = 0; i < count; ++i) { - auto offset = seg_offsets[i]; - dst[i] = std::move(T(field->RawAt(offset))); + static_cast(column->ValueAt(offset))); } } template void ChunkedSegmentSealedImpl::bulk_subscript_ptr_impl( - const ChunkedColumnBase* column, + ChunkedColumnBase* column, const int64_t* seg_offsets, int64_t count, google::protobuf::RepeatedPtrField* dst) { - auto field = reinterpret_cast*>(column); + auto field = reinterpret_cast*>(column); for (int64_t i = 0; i < count; ++i) { auto offset = seg_offsets[i]; dst->at(i) = std::move(T(field->RawAt(offset))); @@ -1401,11 +918,11 @@ ChunkedSegmentSealedImpl::bulk_subscript_ptr_impl( template void ChunkedSegmentSealedImpl::bulk_subscript_array_impl( - const ChunkedColumnBase* column, + ChunkedColumnBase* column, const int64_t* seg_offsets, int64_t count, google::protobuf::RepeatedPtrField* dst) { - auto field = reinterpret_cast(column); + auto field = reinterpret_cast(column); for (int64_t i = 0; i < count; ++i) { auto offset = seg_offsets[i]; dst->at(i) = std::move(field->RawAt(offset)); @@ -1415,14 +932,15 @@ ChunkedSegmentSealedImpl::bulk_subscript_array_impl( // for dense vector void ChunkedSegmentSealedImpl::bulk_subscript_impl(int64_t element_sizeof, - const ChunkedColumnBase* field, + ChunkedColumnBase* field, const int64_t* seg_offsets, int64_t count, void* dst_raw) { auto dst_vec = reinterpret_cast(dst_raw); + auto column = reinterpret_cast(field); for (int64_t i = 0; i < count; ++i) { auto offset = seg_offsets[i]; - auto src = field->ValueAt(offset); + auto src = column->ValueAt(offset); auto dst = dst_vec + i * element_sizeof; memcpy(dst, src, element_sizeof); } @@ -1439,21 +957,11 @@ ChunkedSegmentSealedImpl::ClearData() { num_rows_ = std::nullopt; scalar_indexings_.clear(); vector_indexings_.clear(); - insert_record_.clear(); + insert_record_slot_ = nullptr; fields_.clear(); variable_fields_avg_size_.clear(); stats_.mem_size = 0; } - auto cc = storage::MmapManager::GetInstance().GetChunkCache(); - if (cc == nullptr) { - return; - } - // munmap and remove binlog from chunk cache - for (const auto& iter : field_data_info_.field_infos) { - for (const auto& binlog : iter.second.insert_files) { - cc->Remove(binlog); - } - } } std::unique_ptr @@ -1712,12 +1220,13 @@ ChunkedSegmentSealedImpl::get_raw_data(FieldId field_id, } case DataType::VECTOR_SPARSE_FLOAT: { auto dst = ret->mutable_vectors()->mutable_sparse_float_vector(); + auto col = reinterpret_cast(column.get()); SparseRowsToProto( [&](size_t i) { auto offset = seg_offsets[i]; auto row = static_cast*>( - static_cast(column->ValueAt(offset))); + static_cast(col->ValueAt(offset))); return offset != INVALID_SEG_OFFSET ? row : nullptr; }, count, @@ -1753,38 +1262,40 @@ ChunkedSegmentSealedImpl::bulk_subscript(FieldId field_id, return fill_with_empty(field_id, count); } - if (HasIndex(field_id)) { - // if field has load scalar index, reverse raw data from index - if (!IsVectorDataType(field_meta.get_data_type())) { - // AssertInfo(num_chunk() == 1, - // "num chunk not equal to 1 for sealed segment"); - auto index = chunk_index_impl(field_id, 0); - if (index->HasRawData()) { - return ReverseDataFromIndex( - index, seg_offsets, count, field_meta); - } - return get_raw_data(field_id, field_meta, seg_offsets, count); - } - - std::chrono::high_resolution_clock::time_point get_vector_start = - std::chrono::high_resolution_clock::now(); - - auto vector = get_vector(field_id, seg_offsets, count); - - std::chrono::high_resolution_clock::time_point get_vector_end = - std::chrono::high_resolution_clock::now(); - double get_vector_cost = std::chrono::duration( - get_vector_end - get_vector_start) - .count(); - monitor::internal_core_get_vector_latency.Observe(get_vector_cost / - 1000); - - return vector; + if (!HasIndex(field_id)) { + Assert(get_bit(field_data_ready_bitset_, field_id)); + return get_raw_data(field_id, field_meta, seg_offsets, count); } - Assert(get_bit(field_data_ready_bitset_, field_id)); + auto index_has_raw = HasRawData(field_id.get()); - return get_raw_data(field_id, field_meta, seg_offsets, count); + if (!IsVectorDataType(field_meta.get_data_type())) { + // if field has load scalar index, reverse raw data from index + if (index_has_raw) { + auto index = chunk_index_impl(field_id, 0); + return ReverseDataFromIndex(index, seg_offsets, count, field_meta); + } + return get_raw_data(field_id, field_meta, seg_offsets, count); + } + + std::chrono::high_resolution_clock::time_point get_vector_start = + std::chrono::high_resolution_clock::now(); + + std::unique_ptr vector{nullptr}; + if (index_has_raw) { + vector = get_vector(field_id, seg_offsets, count); + } else { + vector = get_raw_data(field_id, field_meta, seg_offsets, count); + } + + std::chrono::high_resolution_clock::time_point get_vector_end = + std::chrono::high_resolution_clock::now(); + double get_vector_cost = std::chrono::duration( + get_vector_end - get_vector_start) + .count(); + monitor::internal_core_get_vector_latency.Observe(get_vector_cost / 1000); + + return vector; } std::unique_ptr @@ -1808,8 +1319,7 @@ ChunkedSegmentSealedImpl::bulk_subscript( } } auto dst = ret->mutable_scalars()->mutable_json_data()->mutable_data(); - auto field = - reinterpret_cast*>(column.get()); + auto field = reinterpret_cast*>(column.get()); for (int64_t i = 0; i < count; ++i) { auto offset = seg_offsets[i]; dst->at(i) = ExtractSubJson(std::string(field->RawAt(offset)), @@ -1884,7 +1394,8 @@ ChunkedSegmentSealedImpl::search_ids(const IdArray& id_array, for (auto& pk : pks) { std::vector pk_offsets; if (!is_sorted_by_pk_) { - pk_offsets = insert_record_.search_pk(pk, timestamp); + auto ir_accessor = pin_insert_record(); + pk_offsets = ir_accessor->get_cell_of(0)->search_pk(pk, timestamp); } else { pk_offsets = search_pk(pk, timestamp); } @@ -1912,8 +1423,7 @@ ChunkedSegmentSealedImpl::search_ids(const IdArray& id_array, } SegcoreError -ChunkedSegmentSealedImpl::Delete(int64_t reserved_offset, // deprecated - int64_t size, +ChunkedSegmentSealedImpl::Delete(int64_t size, const IdArray* ids, const Timestamp* timestamps_raw) { auto field_id = schema_->get_primary_field_id().value_or(FieldId(-1)); @@ -1927,15 +1437,17 @@ ChunkedSegmentSealedImpl::Delete(int64_t reserved_offset, // deprecated for (int i = 0; i < size; i++) { ordering[i] = std::make_tuple(timestamps_raw[i], pks[i]); } - // if insert_record_ is empty (may be only-load meta but not data for lru-cache at go side), + // if insert record is empty (may be only-load meta but not data for lru-cache at go side), // filtering may cause the deletion lost, skip the filtering to avoid it. - if (!insert_record_.empty_pks()) { - auto end = std::remove_if( - ordering.begin(), - ordering.end(), - [&](const std::tuple& record) { - return !insert_record_.contain(std::get<1>(record)); - }); + auto ir_accessor = pin_insert_record(); + if (!ir_accessor->get_cell_of(0)->empty_pks()) { + auto end = + std::remove_if(ordering.begin(), + ordering.end(), + [&](const std::tuple& record) { + return !ir_accessor->get_cell_of(0)->contain( + std::get<1>(record)); + }); size = end - ordering.begin(); ordering.resize(size); } @@ -1954,7 +1466,7 @@ ChunkedSegmentSealedImpl::Delete(int64_t reserved_offset, // deprecated sort_pks[i] = pk; } - deleted_record_.StreamPush(sort_pks, sort_timestamps.data()); + deleted_record_->StreamPush(sort_pks, sort_timestamps.data()); return SegcoreError::success(); } @@ -1974,7 +1486,9 @@ ChunkedSegmentSealedImpl::LoadSegmentMeta( for (auto& info : segment_meta.metas()) { slice_lengths.push_back(info.row_count()); } - insert_record_.timestamp_index_.set_length_meta(std::move(slice_lengths)); + auto ir_accessor = pin_insert_record(); + ir_accessor->get_cell_of(0)->timestamp_index_.set_length_meta( + std::move(slice_lengths)); PanicInfo(NotImplemented, "unimplemented"); } @@ -1988,19 +1502,21 @@ void ChunkedSegmentSealedImpl::mask_with_timestamps(BitsetTypeView& bitset_chunk, Timestamp timestamp) const { // TODO change the - AssertInfo(insert_record_.timestamps_.num_chunk() == 1, + auto ir_accessor = pin_insert_record(); + auto ir = ir_accessor->get_cell_of(0); + AssertInfo(ir->timestamps().num_chunk() == 1, "num chunk not equal to 1 for sealed segment"); auto timestamps_data = - (const milvus::Timestamp*)insert_record_.timestamps_.get_chunk_data(0); - auto timestamps_data_size = insert_record_.timestamps_.get_chunk_size(0); + (const milvus::Timestamp*)ir->timestamps().get_chunk_data(0); + auto timestamps_data_size = ir->timestamps().get_chunk_size(0); AssertInfo(timestamps_data_size == get_row_count(), fmt::format("Timestamp size not equal to row count: {}, {}", timestamps_data_size, get_row_count())); - auto range = insert_record_.timestamp_index_.get_active_range(timestamp); + auto range = ir->timestamp_index_.get_active_range(timestamp); - // range == (size_, size_) and size_ is this->timestamps_.size(). + // range == (size_, size_) and size_ is this->timestamps().size(). // it means these data are all useful, we don't need to update bitset_chunk. // It can be thought of as an OR operation with another bitmask that is all 0s, but it is not necessary to do so. if (range.first == range.second && range.first == timestamps_data_size) { @@ -2075,18 +1591,16 @@ ChunkedSegmentSealedImpl::generate_interim_index(const FieldId field_id) { segcore_config_, SegmentType::Sealed, is_sparse)); - if (row_count < field_binlog_config->GetBuildThreshold()) { - return false; - } + // if (row_count < field_binlog_config->GetBuildThreshold()) { + // return false; + // } std::shared_ptr vec_data{}; { std::shared_lock lck(mutex_); vec_data = fields_.at(field_id); } - auto dim = - is_sparse - ? dynamic_cast(vec_data.get())->Dim() - : field_meta.get_dim(); + auto dim = is_sparse ? std::numeric_limits::max() + : field_meta.get_dim(); auto build_config = field_binlog_config->GetBuildBaseParams(); build_config[knowhere::meta::DIM] = std::to_string(dim); @@ -2099,8 +1613,10 @@ ChunkedSegmentSealedImpl::generate_interim_index(const FieldId field_id) { knowhere::Version::GetCurrentVersion().VersionNumber()); auto num_chunk = vec_data->num_chunks(); for (int i = 0; i < num_chunk; ++i) { + auto pw = vec_data->GetChunk(i); + auto chunk = pw.get(); auto dataset = knowhere::GenDataSet( - vec_data->chunk_row_nums(i), dim, vec_data->Data(i)); + vec_data->chunk_row_nums(i), dim, chunk->Data()); dataset->SetIsOwner(false); dataset->SetIsSparse(is_sparse); @@ -2131,18 +1647,6 @@ ChunkedSegmentSealedImpl::generate_interim_index(const FieldId field_id) { } void ChunkedSegmentSealedImpl::RemoveFieldFile(const FieldId field_id) { - auto cc = storage::MmapManager::GetInstance().GetChunkCache(); - if (cc == nullptr) { - return; - } - for (const auto& iter : field_data_info_.field_infos) { - if (iter.second.field_id == field_id.get()) { - for (const auto& binlog : iter.second.insert_files) { - cc->Remove(binlog); - } - return; - } - } } void @@ -2162,7 +1666,12 @@ ChunkedSegmentSealedImpl::reopen(SchemaPtr sch) { auto absent_fields = sch->absent_fields(*schema_); for (const auto& field_meta : *absent_fields) { - fill_empty_field(field_meta); + // vector field is not supported to be "added field", thus if a vector + // field is absent, it means for some reason we want to skip loading this + // field. + if (!IsVectorDataType(field_meta.get_data_type())) { + fill_empty_field(field_meta); + } } schema_ = sch; @@ -2170,6 +1679,7 @@ ChunkedSegmentSealedImpl::reopen(SchemaPtr sch) { void ChunkedSegmentSealedImpl::finish_load() { + std::unique_lock lck(mutex_); for (const auto& [field_id, field_meta] : schema_->get_fields()) { if (field_id.get() < START_USER_FIELDID) { continue; @@ -2178,7 +1688,12 @@ ChunkedSegmentSealedImpl::finish_load() { // this shall check the ready bitset here if (!get_bit(field_data_ready_bitset_, field_id) && !get_bit(index_ready_bitset_, field_id)) { - fill_empty_field(field_meta); + // vector field is not supported to be "added field", thus if a vector + // field is absent, it means for some reason we want to skip loading this + // field. + if (!IsVectorDataType(field_meta.get_data_type())) { + fill_empty_field(field_meta); + } } } } @@ -2187,25 +1702,37 @@ void ChunkedSegmentSealedImpl::fill_empty_field(const FieldMeta& field_meta) { int64_t size = num_rows_.value(); AssertInfo(size > 0, "Chunked Sealed segment must have more than 0 row"); - auto column = std::make_shared(field_meta); - auto builder = storage::CreateArrowBuilder(field_meta.get_data_type()); - auto ast = builder->AppendNulls(size); - AssertInfo( - ast.ok(), "append nulls to arrow builder failed: {}", ast.ToString()); - arrow::ArrayVector array_vec; - array_vec.emplace_back(builder->Finish().ValueOrDie()); - - auto chunk = create_chunk( - field_meta, - IsVectorDataType(field_meta.get_data_type()) && - !IsSparseFloatVectorDataType(field_meta.get_data_type()) - ? field_meta.get_dim() - : 1, - array_vec); - column->AddChunk(chunk); + auto field_data_info = FieldDataInfo(field_meta.get_id().get(), size, ""); + std::unique_ptr> translator = + std::make_unique( + get_segment_id(), field_meta, field_data_info, false); + std::shared_ptr column{}; + switch (field_meta.get_data_type()) { + case milvus::DataType::STRING: + case milvus::DataType::VARCHAR: + case milvus::DataType::TEXT: { + column = std::make_shared>( + std::move(translator), field_meta); + break; + } + case milvus::DataType::JSON: { + column = std::make_shared>( + std::move(translator), field_meta); + break; + } + case milvus::DataType::ARRAY: { + column = std::make_shared(std::move(translator), + field_meta); + break; + } + default: { + column = std::make_shared(std::move(translator), + field_meta); + break; + } + } auto field_id = field_meta.get_id(); fields_.emplace(field_id, column); - set_bit(field_data_ready_bitset_, field_id, true); } diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.h b/internal/core/src/segcore/ChunkedSegmentSealedImpl.h index b89f40714f..12f1670b06 100644 --- a/internal/core/src/segcore/ChunkedSegmentSealedImpl.h +++ b/internal/core/src/segcore/ChunkedSegmentSealedImpl.h @@ -14,30 +14,36 @@ #include #include -#include -#include #include #include #include #include #include +#include "folly/executors/InlineExecutor.h" + #include "ConcurrentVector.h" #include "DeletedRecord.h" #include "SealedIndexingRecord.h" #include "SegmentSealed.h" -#include "TimestampIndex.h" #include "common/EasyAssert.h" #include "common/Schema.h" #include "google/protobuf/message_lite.h" #include "mmap/ChunkedColumn.h" -#include "index/ScalarIndex.h" -#include "sys/mman.h" +#include "mmap/Types.h" #include "common/Types.h" #include "common/IndexMeta.h" +#include "cachinglayer/CacheSlot.h" +#include "cachinglayer/CacheSlot.h" namespace milvus::segcore { +namespace storagev1translator { +class InsertRecordTranslator; +} + +using namespace milvus::cachinglayer; + class ChunkedSegmentSealedImpl : public SegmentSealed { public: explicit ChunkedSegmentSealedImpl(SchemaPtr schema, @@ -67,13 +73,9 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { bool Contain(const PkType& pk) const override { - return insert_record_.contain(pk); + return pin_insert_record()->get_cell_of(0)->contain(pk); } - void - LoadFieldData(FieldId field_id, FieldDataInfo& data) override; - void - MapFieldData(const FieldId field_id, FieldDataInfo& data) override; void AddFieldDataInfoForSealed( const LoadFieldDataInfo& field_data_info) override; @@ -117,11 +119,17 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { return iter->second.get(); } - std::pair + // TODO(tiered storage 1): should return a PinWrapper + std::pair GetJsonData(FieldId field_id, size_t offset) const override { - auto column = fields_.at(field_id); + auto column = + std::dynamic_pointer_cast>( + fields_.at(field_id)); bool is_valid = column->IsValid(offset); - return std::make_pair(std::move(column->RawAt(offset)), is_valid); + if (!is_valid) { + return std::make_pair(milvus::Json(), false); + } + return std::make_pair(milvus::Json(column->RawAt(offset)), is_valid); } void @@ -136,12 +144,7 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { public: size_t GetMemoryUsageInBytes() const override { - return stats_.mem_size.load() + deleted_record_.mem_size(); - } - - InsertRecord& - get_insert_record() override { - return insert_record_; + return stats_.mem_size.load() + deleted_record_->mem_size(); } int64_t @@ -156,9 +159,6 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { std::vector search_pk(const PkType& pk, Timestamp timestamp) const override; - std::vector - search_pk(const PkType& pk, int64_t insert_barrier) const override; - template std::vector search_sorted_pk(const PkType& pk, Condition condition) const; @@ -170,11 +170,8 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { bool is_nullable(FieldId field_id) const override { - auto it = fields_.find(field_id); - AssertInfo(it != fields_.end(), - "Cannot find field with field_id: " + - std::to_string(field_id.get())); - return it->second->IsNullable(); + auto& field_meta = schema_->operator[](field_id); + return field_meta.is_nullable(); }; bool @@ -210,8 +207,7 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { debug() const override; SegcoreError - Delete(int64_t reserved_offset, - int64_t size, + Delete(int64_t size, const IdArray* pks, const Timestamp* timestamps) override; @@ -246,32 +242,26 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { protected: // blob and row_count - SpanBase + PinWrapper chunk_data_impl(FieldId field_id, int64_t chunk_id) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_string_view_impl( FieldId field_id, int64_t chunk_id, std::optional> offset_len) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_array_view_impl( FieldId field_id, int64_t chunk_id, std::optional> offset_len) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_view_by_offsets(FieldId field_id, int64_t chunk_id, const FixedVector& offsets) const override; - std::pair> - get_chunk_buffer(FieldId field_id, - int64_t chunk_id, - int64_t start_offset, - int64_t length) const override; - const index::IndexBase* chunk_index_impl(FieldId field_id, int64_t chunk_id) const override; @@ -291,10 +281,13 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { const ConcurrentVector& get_timestamps() const override { - return insert_record_.timestamps_; + return pin_insert_record()->get_cell_of(0)->timestamps_; } private: + void + LoadSystemFieldInternal(FieldId field_id, FieldDataInfo& data); + template static void bulk_subscript_impl(const void* src_raw, @@ -304,35 +297,28 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { template static void - bulk_subscript_impl(const ChunkedColumnBase* field, + bulk_subscript_impl(ChunkedColumnBase* field, const int64_t* seg_offsets, int64_t count, T* dst_raw); template static void - bulk_subscript_impl(const ChunkedColumnBase* field, - const int64_t* seg_offsets, - int64_t count, - void* dst_raw); - - template - static void - bulk_subscript_ptr_impl(const ChunkedColumnBase* field, + bulk_subscript_ptr_impl(ChunkedColumnBase* field, const int64_t* seg_offsets, int64_t count, google::protobuf::RepeatedPtrField* dst_raw); template static void - bulk_subscript_array_impl(const ChunkedColumnBase* column, + bulk_subscript_array_impl(ChunkedColumnBase* column, const int64_t* seg_offsets, int64_t count, google::protobuf::RepeatedPtrField* dst); static void bulk_subscript_impl(int64_t element_sizeof, - const ChunkedColumnBase* field, + ChunkedColumnBase* field, const int64_t* seg_offsets, int64_t count, void* dst_raw); @@ -348,12 +334,8 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { void update_row_count(int64_t row_count) { - // if (row_count_opt_.has_value()) { - // AssertInfo(row_count_opt_.value() == row_count, "load data has different row count from other columns"); - // } else { num_rows_ = row_count; - // } - deleted_record_.set_sealed_row_count(row_count); + deleted_record_->set_sealed_row_count(row_count); } void @@ -375,31 +357,34 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { bool is_system_field_ready() const { - return system_ready_count_ == 2; + return system_ready_count_ == 1; } std::pair, std::vector> search_ids(const IdArray& id_array, Timestamp timestamp) const override; - std::tuple - GetFieldDataPath(FieldId field_id, int64_t offset) const; - void LoadVecIndex(const LoadIndexInfo& info); void LoadScalarIndex(const LoadIndexInfo& info); - void - WarmupChunkCache(const FieldId field_id, bool mmap_enabled) override; - bool generate_interim_index(const FieldId field_id); void fill_empty_field(const FieldMeta& field_meta); + // used only by unit test + std::shared_ptr>> + get_insert_record_slot() const override { + return insert_record_slot_; + } + private: + // InsertRecord needs to pin pk column. + friend class storagev1translator::InsertRecordTranslator; + // mmap descriptor, used in chunk cache storage::MmapChunkDescriptorPtr mmap_descriptor_ = nullptr; // segment loading state @@ -417,17 +402,25 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { // vector field index SealedIndexingRecord vector_indexings_; + std::shared_ptr>> + pin_insert_record() const { + return insert_record_slot_->PinCells({0}) + .via(&folly::InlineExecutor::instance()) + .get(); + } + // inserted fields data and row_ids, timestamps - InsertRecord insert_record_; + mutable std::shared_ptr>> insert_record_slot_; // deleted pks - mutable DeletedRecord deleted_record_; + mutable std::unique_ptr> deleted_record_; LoadFieldDataInfo field_data_info_; SchemaPtr schema_; int64_t id_; - std::unordered_map> fields_; + mutable std::unordered_map> + fields_; std::unordered_set mmap_fields_; // only useful in binlog @@ -443,6 +436,7 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { bool TEST_skip_index_for_retrieve_ = false; // whether the segment is sorted by the pk + // 1. will skip index loading for primary key field bool is_sorted_by_pk_ = false; // used for json expr optimization std::unordered_map; -template -class ThreadSafeVector { - public: - template - void - emplace_to_at_least(int64_t size, Args... args) { - std::lock_guard lck(mutex_); - if (size <= size_) { - return; - } - while (vec_.size() < size) { - vec_.emplace_back(std::forward(args...)); - ++size_; - } - } - const Type& - operator[](int64_t index) const { - std::shared_lock lck(mutex_); - AssertInfo(index < size_, - fmt::format( - "index out of range, index={}, size_={}", index, size_)); - return vec_[index]; - } - - Type& - operator[](int64_t index) { - std::shared_lock lck(mutex_); - AssertInfo(index < size_, - fmt::format( - "index out of range, index={}, size_={}", index, size_)); - return vec_[index]; - } - - int64_t - size() const { - std::shared_lock lck(mutex_); - return size_; - } - - void - clear() { - std::lock_guard lck(mutex_); - size_ = 0; - vec_.clear(); - } - - private: - int64_t size_ = 0; - std::deque vec_; - mutable std::shared_mutex mutex_; -}; - class VectorBase { public: explicit VectorBase(int64_t size_per_chunk) diff --git a/internal/core/src/segcore/DeletedRecord.h b/internal/core/src/segcore/DeletedRecord.h index 9b804dd2cb..2062d50739 100644 --- a/internal/core/src/segcore/DeletedRecord.h +++ b/internal/core/src/segcore/DeletedRecord.h @@ -19,13 +19,9 @@ #include #include -#include "AckResponder.h" -#include "common/Schema.h" +#include "cachinglayer/CacheSlot.h" #include "common/Types.h" -#include "segcore/Record.h" #include "segcore/InsertRecord.h" -#include "segcore/SegmentInterface.h" -#include "ConcurrentVector.h" namespace milvus::segcore { @@ -52,7 +48,12 @@ static int32_t DELETE_PAIR_SIZE = sizeof(std::pair); template class DeletedRecord { public: - DeletedRecord(InsertRecord* insert_record, + using InsertRecordType = std::conditional_t< + is_sealed, + std::shared_ptr>>, + InsertRecord*>; + + DeletedRecord(InsertRecordType insert_record, std::function( const PkType& pk, Timestamp timestamp)> search_pk_func, int64_t segment_id) @@ -108,6 +109,18 @@ class DeletedRecord { Timestamp max_timestamp = 0; SortedDeleteList::Accessor accessor(deleted_lists_); + InsertRecord* insert_record = nullptr; + std::shared_ptr< + milvus::cachinglayer::CellAccessor>> + cell_accessor{nullptr}; + if constexpr (is_sealed) { + cell_accessor = insert_record_->PinCells({0}) + .via(&folly::InlineExecutor::instance()) + .get(); + insert_record = cell_accessor->get_cell_of(0); + } else { + insert_record = insert_record_; + } for (size_t i = 0; i < pks.size(); ++i) { auto deleted_pk = pks[i]; auto deleted_ts = timestamps[i]; @@ -124,7 +137,7 @@ class DeletedRecord { } // if insert record and delete record is same timestamp, // delete not take effect on this record. - if (deleted_ts == insert_record_->timestamps_[row_id]) { + if (deleted_ts == insert_record->timestamps_[row_id]) { continue; } accessor.insert(std::make_pair(deleted_ts, row_id)); @@ -133,7 +146,7 @@ class DeletedRecord { deleted_mask_.set(row_id); } else { // need to add mask size firstly for growing segment - deleted_mask_.resize(insert_record_->size()); + deleted_mask_.resize(insert_record->size()); deleted_mask_.set(row_id); } removed_num++; @@ -322,7 +335,7 @@ class DeletedRecord { public: std::atomic n_ = 0; std::atomic mem_size_ = 0; - InsertRecord* insert_record_; + InsertRecordType insert_record_; std::function(const PkType& pk, Timestamp timestamp)> search_pk_func_; int64_t segment_id_{0}; diff --git a/internal/core/src/segcore/FieldIndexing.h b/internal/core/src/segcore/FieldIndexing.h index 8752e63e2c..7c8dc24fdc 100644 --- a/internal/core/src/segcore/FieldIndexing.h +++ b/internal/core/src/segcore/FieldIndexing.h @@ -20,7 +20,6 @@ #include #include -#include "AckResponder.h" #include "InsertRecord.h" #include "common/FieldMeta.h" #include "common/Schema.h" @@ -29,6 +28,7 @@ #include "knowhere/config.h" #include "log/Log.h" #include "segcore/SegcoreConfig.h" +#include "segcore/InsertRecord.h" #include "index/VectorIndex.h" namespace milvus::segcore { @@ -306,13 +306,12 @@ class IndexingRecord { } // concurrent, reentrant - template void AppendingIndex(int64_t reserved_offset, int64_t size, FieldId fieldId, const DataArray* stream_data, - const InsertRecord& record) { + const InsertRecord& record) { if (!is_in(fieldId)) { return; } @@ -339,13 +338,12 @@ class IndexingRecord { } // concurrent, reentrant - template void AppendingIndex(int64_t reserved_offset, int64_t size, FieldId fieldId, const FieldDataPtr data, - const InsertRecord& record) { + const InsertRecord& record) { if (!is_in(fieldId)) { return; } diff --git a/internal/core/src/segcore/InsertRecord.h b/internal/core/src/segcore/InsertRecord.h index e20db83f7d..425f65cc19 100644 --- a/internal/core/src/segcore/InsertRecord.h +++ b/internal/core/src/segcore/InsertRecord.h @@ -20,19 +20,14 @@ #include #include #include -#include #include "TimestampIndex.h" #include "common/EasyAssert.h" #include "common/Schema.h" #include "common/Types.h" -#include "fmt/format.h" #include "mmap/ChunkedColumn.h" -#include "mmap/Column.h" #include "segcore/AckResponder.h" #include "segcore/ConcurrentVector.h" -#include "segcore/Record.h" -#include "storage/MmapManager.h" namespace milvus::segcore { @@ -279,39 +274,35 @@ class OffsetOrderedArray : public OffsetMap { std::vector> array_; }; -template +template struct InsertRecord { + public: InsertRecord( const Schema& schema, const int64_t size_per_chunk, - const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr) + const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr, + bool called_from_subclass = false) : timestamps_(size_per_chunk), mmap_descriptor_(mmap_descriptor) { + if (called_from_subclass) { + return; + } std::optional pk_field_id = schema.get_primary_field_id(); - + // for sealed segment, only pk field is added. for (auto& field : schema) { auto field_id = field.first; auto& field_meta = field.second; - if (pk2offset_ == nullptr && pk_field_id.has_value() && - pk_field_id.value() == field_id) { + if (pk_field_id.has_value() && pk_field_id.value() == field_id) { + AssertInfo(!field_meta.is_nullable(), + "Primary key should not be nullable"); switch (field_meta.get_data_type()) { case DataType::INT64: { - if constexpr (is_sealed) { - pk2offset_ = - std::make_unique>(); - } else { - pk2offset_ = - std::make_unique>(); - } + pk2offset_ = + std::make_unique>(); break; } case DataType::VARCHAR: { - if constexpr (is_sealed) { - pk2offset_ = std::make_unique< - OffsetOrderedArray>(); - } else { - pk2offset_ = std::make_unique< - OffsetOrderedMap>(); - } + pk2offset_ = + std::make_unique>(); break; } default: { @@ -321,7 +312,6 @@ struct InsertRecord { } } } - append_field_meta(field_id, field_meta, size_per_chunk); } } @@ -344,18 +334,16 @@ struct InsertRecord { } void - insert_pks(milvus::DataType data_type, - const std::shared_ptr& data) { + insert_pks(milvus::DataType data_type, ChunkedColumnBase* data) { std::lock_guard lck(shared_mutex_); int64_t offset = 0; switch (data_type) { case DataType::INT64: { - auto column = std::dynamic_pointer_cast(data); - auto num_chunk = column->num_chunks(); + auto num_chunk = data->num_chunks(); for (int i = 0; i < num_chunk; ++i) { - auto pks = - reinterpret_cast(column->Data(i)); - auto chunk_num_rows = column->chunk_row_nums(i); + auto pw = data->DataOfChunk(i); + auto pks = reinterpret_cast(pw.get()); + auto chunk_num_rows = data->chunk_row_nums(i); for (int j = 0; j < chunk_num_rows; ++j) { pk2offset_->insert(pks[j], offset++); } @@ -363,12 +351,10 @@ struct InsertRecord { break; } case DataType::VARCHAR: { - auto column = std::dynamic_pointer_cast< - ChunkedVariableColumn>(data); - - auto num_chunk = column->num_chunks(); + auto num_chunk = data->num_chunks(); for (int i = 0; i < num_chunk; ++i) { - auto pks = column->StringViews(i).first; + auto pw = data->StringViews(i); + auto pks = pw.get().first; for (auto& pk : pks) { pk2offset_->insert(std::string(pk), offset++); } @@ -383,6 +369,90 @@ struct InsertRecord { } } + bool + empty_pks() const { + std::shared_lock lck(shared_mutex_); + return pk2offset_->empty(); + } + + void + seal_pks() { + std::lock_guard lck(shared_mutex_); + pk2offset_->seal(); + } + + const ConcurrentVector& + timestamps() const { + return timestamps_; + } + + virtual void + clear() { + timestamps_.clear(); + timestamp_index_ = TimestampIndex(); + pk2offset_->clear(); + reserved = 0; + } + + size_t + CellByteSize() const { + return 0; + } + + public: + ConcurrentVector timestamps_; + + std::atomic reserved = 0; + + // used for timestamps index of sealed segment + TimestampIndex timestamp_index_; + + // pks to row offset + std::unique_ptr pk2offset_; + + protected: + storage::MmapChunkDescriptorPtr mmap_descriptor_; + std::unordered_map> data_{}; + mutable std::shared_mutex shared_mutex_{}; +}; + +template <> +struct InsertRecord : public InsertRecord { + public: + InsertRecord( + const Schema& schema, + const int64_t size_per_chunk, + const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr) + : InsertRecord(schema, size_per_chunk, mmap_descriptor, true) { + std::optional pk_field_id = schema.get_primary_field_id(); + for (auto& field : schema) { + auto field_id = field.first; + auto& field_meta = field.second; + if (pk_field_id.has_value() && pk_field_id.value() == field_id) { + AssertInfo(!field_meta.is_nullable(), + "Primary key should not be nullable"); + switch (field_meta.get_data_type()) { + case DataType::INT64: { + pk2offset_ = + std::make_unique>(); + break; + } + case DataType::VARCHAR: { + pk2offset_ = + std::make_unique>(); + break; + } + default: { + PanicInfo(DataTypeInvalid, + fmt::format("unsupported pk type", + field_meta.get_data_type())); + } + } + } + append_field_meta(field_id, field_meta, size_per_chunk); + } + } + void insert_pks(const std::vector& field_datas) { std::lock_guard lck(shared_mutex_); @@ -416,103 +486,6 @@ struct InsertRecord { } } - std::vector - search_pk(const PkType& pk, int64_t insert_barrier) const { - std::shared_lock lck(shared_mutex_); - std::vector res_offsets; - auto offset_iter = pk2offset_->find(pk); - for (auto offset : offset_iter) { - if (offset < insert_barrier) { - res_offsets.emplace_back(offset); - } - } - return res_offsets; - } - - void - insert_pk(const PkType& pk, int64_t offset) { - std::lock_guard lck(shared_mutex_); - pk2offset_->insert(pk, offset); - } - - bool - empty_pks() const { - std::shared_lock lck(shared_mutex_); - return pk2offset_->empty(); - } - - void - seal_pks() { - std::lock_guard lck(shared_mutex_); - pk2offset_->seal(); - } - - // get data without knowing the type - VectorBase* - get_data_base(FieldId field_id) const { - AssertInfo(data_.find(field_id) != data_.end(), - "Cannot find field_data with field_id: " + - std::to_string(field_id.get())); - AssertInfo(data_.at(field_id) != nullptr, - "data_ at i is null" + std::to_string(field_id.get())); - return data_.at(field_id).get(); - } - - // get field data in given type, const version - template - const ConcurrentVector* - get_data(FieldId field_id) const { - auto base_ptr = get_data_base(field_id); - auto ptr = dynamic_cast*>(base_ptr); - Assert(ptr); - return ptr; - } - - // get field data in given type, non-const version - template - ConcurrentVector* - get_data(FieldId field_id) { - auto base_ptr = get_data_base(field_id); - auto ptr = dynamic_cast*>(base_ptr); - Assert(ptr); - return ptr; - } - - ThreadSafeValidDataPtr - get_valid_data(FieldId field_id) const { - AssertInfo(valid_data_.find(field_id) != valid_data_.end(), - "Cannot find valid_data with field_id: " + - std::to_string(field_id.get())); - AssertInfo(valid_data_.at(field_id) != nullptr, - "valid_data_ at i is null" + std::to_string(field_id.get())); - return valid_data_.at(field_id); - } - - bool - is_data_exist(FieldId field_id) const { - return data_.find(field_id) != data_.end(); - } - - bool - is_valid_data_exist(FieldId field_id) const { - return valid_data_.find(field_id) != valid_data_.end(); - } - - SpanBase - get_span_base(FieldId field_id, int64_t chunk_id) const { - auto data = get_data_base(field_id); - if (is_valid_data_exist(field_id)) { - auto size = data->get_chunk_size(chunk_id); - auto element_offset = data->get_element_offset(chunk_id); - return SpanBase( - data->get_chunk_data(chunk_id), - get_valid_data(field_id)->get_chunk_data(element_offset), - size, - data->get_element_size()); - } - return data->get_span_base(chunk_id); - } - void append_field_meta(FieldId field_id, const FieldMeta& field_meta, @@ -602,18 +575,76 @@ struct InsertRecord { } } - // append a column of scalar or sparse float vector type - template void - append_data(FieldId field_id, int64_t size_per_chunk) { - static_assert(IsScalar || IsSparse); - data_.emplace( - field_id, - std::make_unique>( - size_per_chunk, - mmap_descriptor_, - is_valid_data_exist(field_id) ? get_valid_data(field_id) - : nullptr)); + insert_pk(const PkType& pk, int64_t offset) { + std::lock_guard lck(shared_mutex_); + pk2offset_->insert(pk, offset); + } + + // get data without knowing the type + VectorBase* + get_data_base(FieldId field_id) const { + AssertInfo(data_.find(field_id) != data_.end(), + "Cannot find field_data with field_id: " + + std::to_string(field_id.get())); + AssertInfo(data_.at(field_id) != nullptr, + "data_ at i is null" + std::to_string(field_id.get())); + return data_.at(field_id).get(); + } + + // get field data in given type, const version + template + const ConcurrentVector* + get_data(FieldId field_id) const { + auto base_ptr = get_data_base(field_id); + auto ptr = dynamic_cast*>(base_ptr); + Assert(ptr); + return ptr; + } + + // get field data in given type, non-const version + template + ConcurrentVector* + get_data(FieldId field_id) { + auto base_ptr = get_data_base(field_id); + auto ptr = dynamic_cast*>(base_ptr); + Assert(ptr); + return ptr; + } + + ThreadSafeValidDataPtr + get_valid_data(FieldId field_id) const { + AssertInfo(valid_data_.find(field_id) != valid_data_.end(), + "Cannot find valid_data with field_id: " + + std::to_string(field_id.get())); + AssertInfo(valid_data_.at(field_id) != nullptr, + "valid_data_ at i is null" + std::to_string(field_id.get())); + return valid_data_.at(field_id); + } + + bool + is_data_exist(FieldId field_id) const { + return data_.find(field_id) != data_.end(); + } + + bool + is_valid_data_exist(FieldId field_id) const { + return valid_data_.find(field_id) != valid_data_.end(); + } + + SpanBase + get_span_base(FieldId field_id, int64_t chunk_id) const { + auto data = get_data_base(field_id); + if (is_valid_data_exist(field_id)) { + auto size = data->get_chunk_size(chunk_id); + auto element_offset = data->get_element_offset(chunk_id); + return SpanBase( + data->get_chunk_data(chunk_id), + get_valid_data(field_id)->get_chunk_data(element_offset), + size, + data->get_element_size()); + } + return data->get_span_base(chunk_id); } // append a column of scalar type @@ -633,55 +664,48 @@ struct InsertRecord { dim, size_per_chunk, mmap_descriptor_)); } + // append a column of scalar or sparse float vector type + template + void + append_data(FieldId field_id, int64_t size_per_chunk) { + static_assert(IsScalar || IsSparse); + data_.emplace( + field_id, + std::make_unique>( + size_per_chunk, + mmap_descriptor_, + is_valid_data_exist(field_id) ? get_valid_data(field_id) + : nullptr)); + } + void drop_field_data(FieldId field_id) { data_.erase(field_id); valid_data_.erase(field_id); } - const ConcurrentVector& - timestamps() const { - return timestamps_; - } - int64_t size() const { return ack_responder_.GetAck(); } - void - clear() { - timestamps_.clear(); - reserved = 0; - ack_responder_.clear(); - timestamp_index_ = TimestampIndex(); - pk2offset_->clear(); - data_.clear(); - } - bool empty() const { return pk2offset_->empty(); } + void + clear() override { + InsertRecord::clear(); + data_.clear(); + ack_responder_.clear(); + } public: - ConcurrentVector timestamps_; - // used for preInsert of growing segment - std::atomic reserved = 0; AckResponder ack_responder_; - // used for timestamps index of sealed segment - TimestampIndex timestamp_index_; - - // pks to row offset - std::unique_ptr pk2offset_; - private: - std::unordered_map> data_{}; std::unordered_map valid_data_{}; - mutable std::shared_mutex shared_mutex_{}; - storage::MmapChunkDescriptorPtr mmap_descriptor_; }; } // namespace milvus::segcore diff --git a/internal/core/src/segcore/Record.h b/internal/core/src/segcore/Record.h index 27c3ccbb40..ae659152b1 100644 --- a/internal/core/src/segcore/Record.h +++ b/internal/core/src/segcore/Record.h @@ -11,6 +11,8 @@ #pragma once +#include "common/Types.h" + namespace milvus::segcore { template diff --git a/internal/core/src/segcore/SegmentChunkReader.cpp b/internal/core/src/segcore/SegmentChunkReader.cpp index 38c4433954..df87715db7 100644 --- a/internal/core/src/segcore/SegmentChunkReader.cpp +++ b/internal/core/src/segcore/SegmentChunkReader.cpp @@ -39,20 +39,24 @@ SegmentChunkReader::GetChunkDataAccessor(FieldId field_id, }; } } - auto chunk_info = segment_->chunk_data(field_id, current_chunk_id); + // pw is captured by value, each time we need to access a new chunk, we need to + // pin a new Chunk. + auto pw = segment_->chunk_data(field_id, current_chunk_id); + auto chunk_info = pw.get(); auto chunk_data = chunk_info.data(); auto chunk_valid_data = chunk_info.valid_data(); auto current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); return [=, + pw = std::move(pw), ¤t_chunk_id, ¤t_chunk_pos]() mutable -> const data_access_type { if (current_chunk_pos >= current_chunk_size) { current_chunk_id++; current_chunk_pos = 0; - auto chunk_info = - segment_->chunk_data(field_id, current_chunk_id); - chunk_data = chunk_info.data(); - chunk_valid_data = chunk_info.valid_data(); + // the old chunk will be unpinned, pw will now pin the new chunk. + pw = segment_->chunk_data(field_id, current_chunk_id); + chunk_data = pw.get().data(); + chunk_valid_data = pw.get().valid_data(); current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); } @@ -92,26 +96,28 @@ SegmentChunkReader::GetChunkDataAccessor( !storage::MmapManager::GetInstance() .GetMmapConfig() .growing_enable_mmap) { - auto chunk_info = - segment_->chunk_data(field_id, current_chunk_id); + auto pw = segment_->chunk_data(field_id, current_chunk_id); + auto chunk_info = pw.get(); auto chunk_data = chunk_info.data(); auto chunk_valid_data = chunk_info.valid_data(); auto current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); - return [=, + return [pw = std::move(pw), + this, + field_id, + chunk_data, + chunk_valid_data, + current_chunk_size, + // pw = std::move(pw), ¤t_chunk_id, ¤t_chunk_pos]() mutable -> const data_access_type { if (current_chunk_pos >= current_chunk_size) { current_chunk_id++; current_chunk_pos = 0; - chunk_data = - segment_ - ->chunk_data(field_id, current_chunk_id) - .data(); - chunk_valid_data = - segment_ - ->chunk_data(field_id, current_chunk_id) - .valid_data(); + pw = segment_->chunk_data(field_id, + current_chunk_id); + chunk_data = pw.get().data(); + chunk_valid_data = pw.get().valid_data(); current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); } @@ -122,24 +128,24 @@ SegmentChunkReader::GetChunkDataAccessor( return chunk_data[current_chunk_pos++]; }; } else { - auto chunk_info = + auto pw = segment_->chunk_view(field_id, current_chunk_id); - auto current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); return [=, + pw = std::move(pw), ¤t_chunk_id, ¤t_chunk_pos]() mutable -> const data_access_type { if (current_chunk_pos >= current_chunk_size) { current_chunk_id++; current_chunk_pos = 0; - chunk_info = segment_->chunk_view( - field_id, current_chunk_id); + pw = segment_->chunk_view(field_id, + current_chunk_id); current_chunk_size = segment_->chunk_size(field_id, current_chunk_id); } - auto& chunk_data = chunk_info.first; - auto& chunk_valid_data = chunk_info.second; + auto& chunk_data = pw.get().first; + auto& chunk_valid_data = pw.get().second; if (current_chunk_pos < chunk_valid_data.size() && !chunk_valid_data[current_chunk_pos]) { current_chunk_pos++; @@ -204,10 +210,11 @@ SegmentChunkReader::GetChunkDataAccessor(FieldId field_id, }; } } - auto chunk_data = segment_->chunk_data(field_id, chunk_id).data(); - auto chunk_valid_data = - segment_->chunk_data(field_id, chunk_id).valid_data(); - return [chunk_data, chunk_valid_data](int i) -> const data_access_type { + auto pw = segment_->chunk_data(field_id, chunk_id); + return [pw = std::move(pw)](int i) mutable -> const data_access_type { + auto chunk_info = pw.get(); + auto chunk_data = chunk_info.data(); + auto chunk_valid_data = chunk_info.valid_data(); if (chunk_valid_data && !chunk_valid_data[i]) { return std::nullopt; } @@ -237,22 +244,20 @@ SegmentChunkReader::GetChunkDataAccessor(FieldId field_id, !storage::MmapManager::GetInstance() .GetMmapConfig() .growing_enable_mmap) { - auto chunk_data = - segment_->chunk_data(field_id, chunk_id).data(); - auto chunk_valid_data = - segment_->chunk_data(field_id, chunk_id).valid_data(); - return [chunk_data, chunk_valid_data](int i) -> const data_access_type { + auto pw = segment_->chunk_data(field_id, chunk_id); + return [pw = std::move(pw)](int i) mutable -> const data_access_type { + auto chunk_data = pw.get().data(); + auto chunk_valid_data = pw.get().valid_data(); if (chunk_valid_data && !chunk_valid_data[i]) { return std::nullopt; } return chunk_data[i]; }; } else { - auto chunk_info = - segment_->chunk_view(field_id, chunk_id); - return [chunk_data = std::move(chunk_info.first), - chunk_valid_data = std::move(chunk_info.second)]( - int i) -> const data_access_type { + auto pw = segment_->chunk_view(field_id, chunk_id); + return [pw = std::move(pw)](int i) mutable -> const data_access_type { + auto chunk_data = pw.get().first; + auto chunk_valid_data = pw.get().second; if (i < chunk_valid_data.size() && !chunk_valid_data[i]) { return std::nullopt; } diff --git a/internal/core/src/segcore/SegmentGrowingImpl.cpp b/internal/core/src/segcore/SegmentGrowingImpl.cpp index d52de1c661..994aaed69f 100644 --- a/internal/core/src/segcore/SegmentGrowingImpl.cpp +++ b/internal/core/src/segcore/SegmentGrowingImpl.cpp @@ -22,6 +22,7 @@ #include #include +#include "cachinglayer/CacheSlot.h" #include "common/Consts.h" #include "common/EasyAssert.h" #include "common/FieldData.h" @@ -46,6 +47,8 @@ namespace milvus::segcore { +using namespace milvus::cachinglayer; + int64_t SegmentGrowingImpl::PreInsert(int64_t size) { auto reserved_begin = insert_record_.reserved.fetch_add(size); @@ -509,8 +512,7 @@ SegmentGrowingImpl::load_column_group_data_internal( } SegcoreError -SegmentGrowingImpl::Delete(int64_t reserved_begin, - int64_t size, +SegmentGrowingImpl::Delete(int64_t size, const IdArray* ids, const Timestamp* timestamps_raw) { auto field_id = schema_->get_primary_field_id().value_or(FieldId(-1)); @@ -573,12 +575,13 @@ SegmentGrowingImpl::LoadDeletedRecord(const LoadDeletedRecordInfo& info) { deleted_record_.LoadPush(pks, timestamps); } -SpanBase +PinWrapper SegmentGrowingImpl::chunk_data_impl(FieldId field_id, int64_t chunk_id) const { - return get_insert_record().get_span_base(field_id, chunk_id); + return PinWrapper( + get_insert_record().get_span_base(field_id, chunk_id)); } -std::pair, FixedVector> +PinWrapper, FixedVector>> SegmentGrowingImpl::chunk_string_view_impl( FieldId field_id, int64_t chunk_id, @@ -588,7 +591,7 @@ SegmentGrowingImpl::chunk_string_view_impl( "chunk string view impl not implement for growing segment"); } -std::pair, FixedVector> +PinWrapper, FixedVector>> SegmentGrowingImpl::chunk_array_view_impl( FieldId field_id, int64_t chunk_id, @@ -598,7 +601,7 @@ SegmentGrowingImpl::chunk_array_view_impl( "chunk array view impl not implement for growing segment"); } -std::pair, FixedVector> +PinWrapper, FixedVector>> SegmentGrowingImpl::chunk_view_by_offsets( FieldId field_id, int64_t chunk_id, @@ -1154,7 +1157,7 @@ SegmentGrowingImpl::CreateJSONIndex(FieldId field_id) { json_indexes_[field_id] = std::move(index); } -std::pair +std::pair SegmentGrowingImpl::GetJsonData(FieldId field_id, size_t offset) const { auto vec_ptr = dynamic_cast*>( insert_record_.get_data_base(field_id)); @@ -1162,10 +1165,9 @@ SegmentGrowingImpl::GetJsonData(FieldId field_id, size_t offset) const { auto& field_meta = schema_->operator[](field_id); if (field_meta.is_nullable()) { auto valid_data_ptr = insert_record_.get_valid_data(field_id); - return std::make_pair(std::string_view(src[offset]), - valid_data_ptr->is_valid(offset)); + return std::make_pair(src[offset], valid_data_ptr->is_valid(offset)); } - return std::make_pair(std::string_view(src[offset]), true); + return std::make_pair(src[offset], true); } void diff --git a/internal/core/src/segcore/SegmentGrowingImpl.h b/internal/core/src/segcore/SegmentGrowingImpl.h index 6641a9ba9c..789f753a93 100644 --- a/internal/core/src/segcore/SegmentGrowingImpl.h +++ b/internal/core/src/segcore/SegmentGrowingImpl.h @@ -21,6 +21,7 @@ #include #include +#include "cachinglayer/CacheSlot.h" #include "AckResponder.h" #include "ConcurrentVector.h" #include "DeletedRecord.h" @@ -35,6 +36,8 @@ namespace milvus::segcore { +using namespace milvus::cachinglayer; + class SegmentGrowingImpl : public SegmentGrowing { public: int64_t @@ -54,8 +57,7 @@ class SegmentGrowingImpl : public SegmentGrowing { // TODO: add id into delete log, possibly bitmap SegcoreError - Delete(int64_t reserved_offset, - int64_t size, + Delete(int64_t size, const IdArray* pks, const Timestamp* timestamps) override; @@ -107,7 +109,7 @@ class SegmentGrowingImpl : public SegmentGrowing { finish_load() override; public: - const InsertRecord<>& + const InsertRecord& get_insert_record() const { return insert_record_; } @@ -248,7 +250,7 @@ class SegmentGrowingImpl : public SegmentGrowing { int64_t count, const std::vector& dynamic_field_names) const override; - virtual std::pair + virtual std::pair GetJsonData(FieldId field_id, size_t offset) const override; public: @@ -374,11 +376,6 @@ class SegmentGrowingImpl : public SegmentGrowing { return insert_record_.search_pk(pk, timestamp); } - std::vector - search_pk(const PkType& pk, int64_t insert_barrier) const override { - return insert_record_.search_pk(pk, insert_barrier); - } - bool is_field_exist(FieldId field_id) const override { return schema_->get_fields().find(field_id) != @@ -389,36 +386,26 @@ class SegmentGrowingImpl : public SegmentGrowing { int64_t num_chunk(FieldId field_id) const override; - SpanBase + PinWrapper chunk_data_impl(FieldId field_id, int64_t chunk_id) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_string_view_impl( FieldId field_id, int64_t chunk_id, std::optional> offset_len) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_array_view_impl( FieldId field_id, int64_t chunk_id, std::optional> offset_len) const override; - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_view_by_offsets(FieldId field_id, int64_t chunk_id, const FixedVector& offsets) const override; - std::pair> - get_chunk_buffer(FieldId field_id, - int64_t chunk_id, - int64_t start_offset, - int64_t length) const override { - PanicInfo( - ErrorCode::Unsupported, - "get_chunk_buffer interface not supported for growing segment"); - } - void check_search(const query::Plan* plan) const override { Assert(plan); @@ -478,10 +465,6 @@ class SegmentGrowingImpl : public SegmentGrowing { SegmentStats stats_{}; }; -const static IndexMetaPtr empty_index_meta = - std::make_shared(1024, - std::map()); - inline SegmentGrowingPtr CreateGrowingSegment( SchemaPtr schema, diff --git a/internal/core/src/segcore/SegmentInterface.cpp b/internal/core/src/segcore/SegmentInterface.cpp index 25cbb3f334..cc55d273cc 100644 --- a/internal/core/src/segcore/SegmentInterface.cpp +++ b/internal/core/src/segcore/SegmentInterface.cpp @@ -18,6 +18,7 @@ #include "common/SystemProperty.h" #include "common/Tracer.h" #include "common/Types.h" +#include "monitor/prometheus_client.h" #include "query/ExecPlanNodeVisitor.h" namespace milvus::segcore { diff --git a/internal/core/src/segcore/SegmentInterface.h b/internal/core/src/segcore/SegmentInterface.h index b6d7b6b42f..4e6ede176f 100644 --- a/internal/core/src/segcore/SegmentInterface.h +++ b/internal/core/src/segcore/SegmentInterface.h @@ -20,8 +20,10 @@ #include #include "FieldIndexing.h" +#include "cachinglayer/CacheSlot.h" #include "common/Common.h" #include "common/EasyAssert.h" +#include "common/Json.h" #include "common/Schema.h" #include "common/Span.h" #include "common/SystemProperty.h" @@ -36,12 +38,13 @@ #include "pb/segcore.pb.h" #include "index/IndexInfo.h" #include "index/SkipIndex.h" -#include "mmap/Column.h" #include "index/TextMatchIndex.h" #include "index/JsonKeyStatsInvertedIndex.h" namespace milvus::segcore { +using namespace milvus::cachinglayer; + struct SegmentStats { // we stat the memory size used by the segment, // including the insert data and delete data. @@ -105,14 +108,8 @@ class SegmentInterface { int64_t num_rows, int64_t field_size) = 0; - // virtual int64_t - // PreDelete(int64_t size) = 0; - virtual SegcoreError - Delete(int64_t reserved_offset, - int64_t size, - const IdArray* pks, - const Timestamp* timestamps) = 0; + Delete(int64_t size, const IdArray* pks, const Timestamp* timestamps) = 0; virtual void LoadDeletedRecord(const LoadDeletedRecordInfo& info) = 0; @@ -145,7 +142,7 @@ class SegmentInterface { virtual index::JsonKeyStatsInvertedIndex* GetJsonKeyIndex(FieldId field_id) const = 0; - virtual std::pair + virtual std::pair GetJsonData(FieldId field_id, size_t offset) const = 0; virtual void @@ -166,40 +163,40 @@ class SegmentInterface { class SegmentInternalInterface : public SegmentInterface { public: template - Span + PinWrapper> chunk_data(FieldId field_id, int64_t chunk_id) const { - return static_cast>(chunk_data_impl(field_id, chunk_id)); + return chunk_data_impl(field_id, chunk_id) + .transform>([](SpanBase&& span_base) { + return static_cast>(span_base); + }); } template - std::pair, FixedVector> + PinWrapper, FixedVector>> chunk_view(FieldId field_id, int64_t chunk_id, std::optional> offset_len = std::nullopt) const { if constexpr (std::is_same_v) { - auto [string_views, valid_data] = - chunk_string_view_impl(field_id, chunk_id, offset_len); - return std::make_pair(std::move(string_views), - std::move(valid_data)); + return chunk_string_view_impl(field_id, chunk_id, offset_len); } else if constexpr (std::is_same_v) { - auto [array_views, valid_data] = - chunk_array_view_impl(field_id, chunk_id, offset_len); - return std::make_pair(array_views, valid_data); + return chunk_array_view_impl(field_id, chunk_id, offset_len); } else if constexpr (std::is_same_v) { - auto [string_views, valid_data] = - chunk_string_view_impl(field_id, chunk_id, offset_len); + auto pw = chunk_string_view_impl(field_id, chunk_id, offset_len); + auto [string_views, valid_data] = pw.get(); std::vector res; res.reserve(string_views.size()); for (const auto& str_view : string_views) { res.emplace_back(str_view); } - return {std::move(res), std::move(valid_data)}; + return PinWrapper< + std::pair, FixedVector>>( + {std::move(res), std::move(valid_data)}); } } template - std::pair, FixedVector> + PinWrapper, FixedVector>> get_batch_views(FieldId field_id, int64_t chunk_id, int64_t start_offset, @@ -213,7 +210,7 @@ class SegmentInternalInterface : public SegmentInterface { } template - std::pair, FixedVector> + PinWrapper, FixedVector>> get_views_by_offsets(FieldId field_id, int64_t chunk_id, const FixedVector& offsets) const { @@ -221,16 +218,20 @@ class SegmentInternalInterface : public SegmentInterface { PanicInfo(ErrorCode::Unsupported, "get chunk views not supported for growing segment"); } - auto chunk_view = chunk_view_by_offsets(field_id, chunk_id, offsets); + auto pw = chunk_view_by_offsets(field_id, chunk_id, offsets); if constexpr (std::is_same_v) { - return chunk_view; + return pw; } else { + static_assert(std::is_same_v, + "only Json is supported for get_views_by_offsets"); std::vector res; - res.reserve(chunk_view.first.size()); - for (const auto& view : chunk_view.first) { + res.reserve(pw.get().first.size()); + for (const auto& view : pw.get().first) { res.emplace_back(view); } - return {res, chunk_view.second}; + return PinWrapper< + std::pair, FixedVector>>( + {std::move(res), pw.get().second}); } } @@ -461,30 +462,25 @@ class SegmentInternalInterface : public SegmentInterface { protected: // todo: use an Unified struct for all type in growing/seal segment to store data and valid_data. // internal API: return chunk_data in span - virtual SpanBase + virtual PinWrapper chunk_data_impl(FieldId field_id, int64_t chunk_id) const = 0; // internal API: return chunk string views in vector - virtual std::pair, FixedVector> + virtual PinWrapper< + std::pair, FixedVector>> chunk_string_view_impl(FieldId field_id, int64_t chunk_id, std::optional> offset_len = std::nullopt) const = 0; - virtual std::pair, FixedVector> + virtual PinWrapper, FixedVector>> chunk_array_view_impl(FieldId field_id, int64_t chunk_id, std::optional> offset_len = std::nullopt) const = 0; - // internal API: return buffer reference to field chunk data located from start_offset - virtual std::pair> - get_chunk_buffer(FieldId field_id, - int64_t chunk_id, - int64_t start_offset, - int64_t length) const = 0; - - virtual std::pair, FixedVector> + virtual PinWrapper< + std::pair, FixedVector>> chunk_view_by_offsets(FieldId field_id, int64_t chunk_id, const FixedVector& offsets) const = 0; @@ -531,9 +527,6 @@ class SegmentInternalInterface : public SegmentInterface { virtual std::vector search_pk(const PkType& pk, Timestamp timestamp) const = 0; - virtual std::vector - search_pk(const PkType& pk, int64_t insert_barrier) const = 0; - protected: mutable std::shared_mutex mutex_; // fieldID -> std::pair diff --git a/internal/core/src/segcore/SegmentSealed.h b/internal/core/src/segcore/SegmentSealed.h index 4b6fd8eab3..d2e10dbc41 100644 --- a/internal/core/src/segcore/SegmentSealed.h +++ b/internal/core/src/segcore/SegmentSealed.h @@ -13,13 +13,13 @@ #include #include -#include #include "common/LoadInfo.h" #include "common/Types.h" -#include "index/JsonInvertedIndex.h" #include "pb/segcore.pb.h" +#include "segcore/InsertRecord.h" #include "segcore/SegmentInterface.h" +#include "cachinglayer/CacheSlot.h" #include "segcore/Types.h" namespace milvus::segcore { @@ -35,15 +35,9 @@ class SegmentSealed : public SegmentInternalInterface { virtual void DropFieldData(const FieldId field_id) = 0; - virtual void - LoadFieldData(FieldId field_id, FieldDataInfo& data) = 0; - virtual void - MapFieldData(const FieldId field_id, FieldDataInfo& data) = 0; virtual void AddFieldDataInfoForSealed(const LoadFieldDataInfo& field_data_info) = 0; virtual void - WarmupChunkCache(const FieldId field_id, bool mmap_enabled) = 0; - virtual void RemoveFieldFile(const FieldId field_id) = 0; virtual void ClearData() = 0; @@ -54,9 +48,6 @@ class SegmentSealed : public SegmentInternalInterface { LoadTextIndex(FieldId field_id, std::unique_ptr index) = 0; - virtual InsertRecord& - get_insert_record() = 0; - virtual index::IndexBase* GetJsonIndex(FieldId field_id, std::string path) const override { JSONIndexKey key; @@ -74,12 +65,6 @@ class SegmentSealed : public SegmentInternalInterface { FieldId field_id, std::unique_ptr index) = 0; - virtual index::JsonKeyStatsInvertedIndex* - GetJsonKeyIndex(FieldId field_id) const = 0; - - virtual std::pair - GetJsonData(FieldId field_id, size_t offset) const = 0; - SegmentType type() const override { return SegmentType::Sealed; @@ -119,6 +104,9 @@ class SegmentSealed : public SegmentInternalInterface { index->second->JsonCastType() == DataType::DOUBLE); } + virtual std::shared_ptr>> + get_insert_record_slot() const = 0; + protected: struct JSONIndexKey { FieldId field_id; diff --git a/internal/core/src/segcore/Utils.h b/internal/core/src/segcore/Utils.h index a946fae20b..213ccc5c97 100644 --- a/internal/core/src/segcore/Utils.h +++ b/internal/core/src/segcore/Utils.h @@ -11,23 +11,15 @@ #pragma once -#include -#include #include -#include #include #include -#include #include #include "common/FieldData.h" -#include "common/QueryResult.h" -// #include "common/Schema.h" #include "common/Types.h" #include "index/Index.h" -#include "log/Log.h" -#include "segcore/DeletedRecord.h" -#include "segcore/InsertRecord.h" +#include "segcore/ConcurrentVector.h" namespace milvus::segcore { diff --git a/internal/core/src/segcore/load_field_data_c.cpp b/internal/core/src/segcore/load_field_data_c.cpp index ddab5a28f4..6efea06ccd 100644 --- a/internal/core/src/segcore/load_field_data_c.cpp +++ b/internal/core/src/segcore/load_field_data_c.cpp @@ -89,12 +89,6 @@ AppendMMapDirPath(CLoadFieldDataInfo c_load_field_data_info, load_field_data_info->mmap_dir_path = std::string(c_dir_path); } -void -SetUri(CLoadFieldDataInfo c_load_field_data_info, const char* uri) { - auto load_field_data_info = (LoadFieldDataInfo*)c_load_field_data_info; - load_field_data_info->url = std::string(uri); -} - void SetStorageVersion(CLoadFieldDataInfo c_load_field_data_info, int64_t storage_version) { diff --git a/internal/core/src/segcore/load_field_data_c.h b/internal/core/src/segcore/load_field_data_c.h index 0b3b7dee9f..dc2166f7cf 100644 --- a/internal/core/src/segcore/load_field_data_c.h +++ b/internal/core/src/segcore/load_field_data_c.h @@ -46,9 +46,6 @@ void AppendMMapDirPath(CLoadFieldDataInfo c_load_field_data_info, const char* dir_path); -void -SetUri(CLoadFieldDataInfo c_load_field_data_info, const char* uri); - void SetStorageVersion(CLoadFieldDataInfo c_load_field_data_info, int64_t storage_version); diff --git a/internal/core/src/segcore/reduce/Reduce.cpp b/internal/core/src/segcore/reduce/Reduce.cpp index 6bdc8cc130..558cdc3a72 100644 --- a/internal/core/src/segcore/reduce/Reduce.cpp +++ b/internal/core/src/segcore/reduce/Reduce.cpp @@ -15,9 +15,10 @@ #include #include +#include "common/EasyAssert.h" +#include "monitor/prometheus_client.h" #include "segcore/SegmentInterface.h" #include "segcore/Utils.h" -#include "common/EasyAssert.h" #include "segcore/pkVisitor.h" #include "segcore/ReduceUtils.h" diff --git a/internal/core/src/segcore/segcore_init_c.cpp b/internal/core/src/segcore/segcore_init_c.cpp index dbb74db5ab..eadb96a8bd 100644 --- a/internal/core/src/segcore/segcore_init_c.cpp +++ b/internal/core/src/segcore/segcore_init_c.cpp @@ -15,6 +15,7 @@ #include "log/Log.h" #include "segcore/SegcoreConfig.h" #include "segcore/segcore_init_c.h" +#include "cachinglayer/Manager.h" namespace milvus::segcore { @@ -115,4 +116,12 @@ SetThreadName(const char* name) { #endif } +extern "C" void +ConfigureTieredStorage(const bool enabled_globally, + const int64_t memory_limit_bytes, + const int64_t disk_limit_bytes) { + milvus::cachinglayer::Manager::ConfigureTieredStorage( + enabled_globally, memory_limit_bytes, disk_limit_bytes); +} + } // namespace milvus::segcore diff --git a/internal/core/src/segcore/segcore_init_c.h b/internal/core/src/segcore/segcore_init_c.h index 4a060c2db1..1c1cf0dac2 100644 --- a/internal/core/src/segcore/segcore_init_c.h +++ b/internal/core/src/segcore/segcore_init_c.h @@ -62,6 +62,11 @@ GetMinimalIndexVersion(); void SetThreadName(const char*); +void +ConfigureTieredStorage(const bool enabled_globally, + const int64_t memory_limit_bytes, + const int64_t disk_limit_bytes); + #ifdef __cplusplus } #endif diff --git a/internal/core/src/segcore/segment_c.cpp b/internal/core/src/segcore/segment_c.cpp index 03479137fd..03c623f825 100644 --- a/internal/core/src/segcore/segment_c.cpp +++ b/internal/core/src/segcore/segment_c.cpp @@ -36,6 +36,7 @@ #include "segcore/SegmentSealed.h" #include "segcore/ChunkedSegmentSealedImpl.h" #include "mmap/Types.h" +#include "storage/RemoteChunkManagerSingleton.h" ////////////////////////////// common interfaces ////////////////////////////// CStatus @@ -319,7 +320,6 @@ PreInsert(CSegmentInterface c_segment, int64_t size, int64_t* offset) { CStatus Delete(CSegmentInterface c_segment, - int64_t reserved_offset, // deprecated int64_t size, const uint8_t* ids, const uint64_t ids_size, @@ -329,8 +329,7 @@ Delete(CSegmentInterface c_segment, auto suc = pks->ParseFromArray(ids, ids_size); AssertInfo(suc, "failed to parse pks from ids"); try { - auto res = - segment->Delete(reserved_offset, size, pks.get(), timestamps); + auto res = segment->Delete(size, pks.get(), timestamps); return milvus::SuccessCStatus(); } catch (std::exception& e) { return milvus::FailureCStatus(&e); @@ -353,49 +352,6 @@ LoadFieldData(CSegmentInterface c_segment, } } -// just for test -CStatus -LoadFieldRawData(CSegmentInterface c_segment, - int64_t field_id, - const void* data, - int64_t row_count) { - try { - auto segment_interface = - reinterpret_cast(c_segment); - auto segment = - dynamic_cast(segment_interface); - AssertInfo(segment != nullptr, "segment conversion failed"); - milvus::DataType data_type; - int64_t dim = 1; - if (milvus::SystemProperty::Instance().IsSystem( - milvus::FieldId(field_id))) { - data_type = milvus::DataType::INT64; - } else { - auto field_meta = segment->get_schema()[milvus::FieldId(field_id)]; - data_type = field_meta.get_data_type(); - - if (milvus::IsVectorDataType(data_type) && - !milvus::IsSparseFloatVectorDataType(data_type)) { - dim = field_meta.get_dim(); - } - } - auto field_data = - milvus::storage::CreateFieldData(data_type, false, dim); - field_data->FillFieldData(data, row_count); - auto arrow_data_wrapper = - milvus::storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = milvus::FieldDataInfo{ - field_id, - static_cast(row_count), - std::vector>{ - arrow_data_wrapper}}; - segment->LoadFieldData(milvus::FieldId(field_id), field_data_info); - return milvus::SuccessCStatus(); - } catch (std::exception& e) { - return milvus::FailureCStatus(&e); - } -} - CStatus LoadDeletedRecord(CSegmentInterface c_segment, CLoadDeletedRecordInfo deleted_record_info) { @@ -608,17 +564,7 @@ CStatus WarmupChunkCache(CSegmentInterface c_segment, int64_t field_id, bool mmap_enabled) { - try { - auto segment_interface = - reinterpret_cast(c_segment); - auto segment = - dynamic_cast(segment_interface); - AssertInfo(segment != nullptr, "segment conversion failed"); - segment->WarmupChunkCache(milvus::FieldId(field_id), mmap_enabled); - return milvus::SuccessCStatus(); - } catch (std::exception& e) { - return milvus::FailureCStatus(milvus::UnexpectedError, e.what()); - } + return milvus::SuccessCStatus(); } void diff --git a/internal/core/src/segcore/segment_c.h b/internal/core/src/segcore/segment_c.h index 6c8b610b65..147cf47bcf 100644 --- a/internal/core/src/segcore/segment_c.h +++ b/internal/core/src/segcore/segment_c.h @@ -105,12 +105,6 @@ CStatus LoadFieldData(CSegmentInterface c_segment, CLoadFieldDataInfo load_field_data_info); -CStatus -LoadFieldRawData(CSegmentInterface c_segment, - int64_t field_id, - const void* data, - int64_t row_count); - CStatus LoadDeletedRecord(CSegmentInterface c_segment, CLoadDeletedRecordInfo deleted_record_info); @@ -136,6 +130,8 @@ UpdateFieldRawDataSize(CSegmentInterface c_segment, int64_t num_rows, int64_t field_data_size); +// This function is currently used only in test. +// Current implement supports only dropping of non-system fields. CStatus DropFieldData(CSegmentInterface c_segment, int64_t field_id); @@ -160,7 +156,6 @@ ExistPk(CSegmentInterface c_segment, CStatus Delete(CSegmentInterface c_segment, - int64_t reserved_offset, int64_t size, const uint8_t* ids, const uint64_t ids_size, diff --git a/internal/core/src/segcore/storagev1translator/ChunkTranslator.cpp b/internal/core/src/segcore/storagev1translator/ChunkTranslator.cpp new file mode 100644 index 0000000000..a411eac0c6 --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/ChunkTranslator.cpp @@ -0,0 +1,166 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include "segcore/storagev1translator/ChunkTranslator.h" + +#include +#include +#include +#include + +#include "cachinglayer/Utils.h" +#include "common/ChunkWriter.h" +#include "common/EasyAssert.h" +#include "common/Types.h" +#include "common/SystemProperty.h" +#include "segcore/Utils.h" +#include "storage/ThreadPools.h" +#include "mmap/Types.h" + +namespace milvus::segcore::storagev1translator { + +ChunkTranslator::ChunkTranslator(int64_t segment_id, + FieldMeta field_meta, + FieldDataInfo field_data_info, + std::vector insert_files, + bool use_mmap) + : segment_id_(segment_id), + key_(fmt::format("seg_{}_f_{}", segment_id, field_data_info.field_id)), + use_mmap_(use_mmap), + meta_(use_mmap ? milvus::cachinglayer::StorageType::DISK + : milvus::cachinglayer::StorageType::MEMORY) { + chunks_.resize(insert_files.size()); + AssertInfo( + !SystemProperty::Instance().IsSystem(FieldId(field_data_info.field_id)), + "ChunkTranslator not supported for system field"); + + auto parallel_degree = + static_cast(DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE); + auto& pool = ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE); + pool.Submit(LoadArrowReaderFromRemote, + insert_files, + field_data_info.arrow_reader_channel); + LOG_INFO("segment {} submits load field {} task to thread pool", + segment_id_, + field_data_info.field_id); + + auto data_type = field_meta.get_data_type(); + auto cid = 0; + auto row_count = 0; + meta_.num_rows_until_chunk_.push_back(0); + if (!use_mmap_) { + std::shared_ptr r; + while (field_data_info.arrow_reader_channel->pop(r)) { + arrow::ArrayVector array_vec = + read_single_column_batches(r->reader); + auto chunk = + create_chunk(field_meta, + IsVectorDataType(data_type) && + !IsSparseFloatVectorDataType(data_type) + ? field_meta.get_dim() + : 1, + array_vec) + .release(); + chunks_[cid] = chunk; + row_count += chunk->RowNums(); + meta_.num_rows_until_chunk_.push_back(row_count); + cid++; + } + } else { + auto filepath = std::filesystem::path(field_data_info.mmap_dir_path) / + std::to_string(segment_id_) / + std::to_string(field_data_info.field_id); + auto dir = filepath.parent_path(); + std::filesystem::create_directories(dir); + + auto file = File::Open(filepath.string(), O_CREAT | O_TRUNC | O_RDWR); + + std::shared_ptr r; + size_t file_offset = 0; + std::vector> chunks; + while (field_data_info.arrow_reader_channel->pop(r)) { + arrow::ArrayVector array_vec = + read_single_column_batches(r->reader); + auto chunk = + create_chunk(field_meta, + IsVectorDataType(data_type) && + !IsSparseFloatVectorDataType(data_type) + ? field_meta.get_dim() + : 1, + file, + file_offset, + array_vec) + .release(); + chunks_[cid] = chunk; + row_count += chunk->RowNums(); + meta_.num_rows_until_chunk_.push_back(row_count); + cid++; + file_offset += chunk->Size(); + } + } + AssertInfo(row_count == field_data_info.row_count, + fmt::format("data lost while loading column {}: loaded " + "num rows {} but expected {}", + field_data_info.field_id, + row_count, + field_data_info.row_count)); +} + +ChunkTranslator::~ChunkTranslator() { + for (auto chunk : chunks_) { + if (chunk != nullptr) { + // let the Chunk to be deleted by the unique_ptr + auto chunk_ptr = std::unique_ptr(chunk); + } + } +} + +size_t +ChunkTranslator::num_cells() const { + return chunks_.size(); +} + +milvus::cachinglayer::cid_t +ChunkTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const { + return uid; +} + +milvus::cachinglayer::ResourceUsage +ChunkTranslator::estimated_byte_size_of_cell( + milvus::cachinglayer::cid_t cid) const { + return {0, 0}; +} + +const std::string& +ChunkTranslator::key() const { + return key_; +} + +std::vector< + std::pair>> +ChunkTranslator::get_cells( + const std::vector& cids) { + std::vector< + std::pair>> + cells; + for (auto cid : cids) { + AssertInfo(chunks_[cid] != nullptr, + "ChunkTranslator::get_cells called again on cell {} of " + "CacheSlot {}.", + cid, + key_); + cells.emplace_back(cid, std::unique_ptr(chunks_[cid])); + chunks_[cid] = nullptr; + } + return cells; +} + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/segcore/storagev1translator/ChunkTranslator.h b/internal/core/src/segcore/storagev1translator/ChunkTranslator.h new file mode 100644 index 0000000000..77d62c1a7d --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/ChunkTranslator.h @@ -0,0 +1,71 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" +#include "common/Chunk.h" +#include "mmap/Types.h" + +namespace milvus::segcore::storagev1translator { + +struct CTMeta : public milvus::cachinglayer::Meta { + std::vector num_rows_until_chunk_; + CTMeta(milvus::cachinglayer::StorageType storage_type) + : milvus::cachinglayer::Meta(storage_type) { + } +}; + +// This class will load all cells(Chunks) in ctor, and move them out during get_cells. +// This should be used only in storagev1(no eviction allowed), thus trying to get a +// same cell a second time will result in exception. +// For this translator each Chunk is a CacheCell, cid_t and uid_ is the same. +class ChunkTranslator : public milvus::cachinglayer::Translator { + public: + ChunkTranslator(int64_t segment_id, + FieldMeta field_meta, + FieldDataInfo field_data_info, + std::vector insert_files, + bool use_mmap); + + ~ChunkTranslator() override; + + size_t + num_cells() const override; + milvus::cachinglayer::cid_t + cell_id_of(milvus::cachinglayer::uid_t uid) const override; + milvus::cachinglayer::ResourceUsage + estimated_byte_size_of_cell(milvus::cachinglayer::cid_t cid) const override; + const std::string& + key() const override; + std::vector< + std::pair>> + get_cells(const std::vector& cids) override; + + // TODO: info other than get_cels() should all be in meta() + milvus::cachinglayer::Meta* + meta() override { + return &meta_; + } + + private: + int64_t segment_id_; + std::string key_; + bool use_mmap_; + std::vector chunks_; + CTMeta meta_; +}; + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.cpp b/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.cpp new file mode 100644 index 0000000000..cd2e655417 --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.cpp @@ -0,0 +1,96 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include "segcore/storagev1translator/DefaultValueChunkTranslator.h" + +#include "common/ChunkWriter.h" +#include "storage/Util.h" + +namespace milvus::segcore::storagev1translator { + +DefaultValueChunkTranslator::DefaultValueChunkTranslator( + int64_t segment_id, + FieldMeta field_meta, + FieldDataInfo field_data_info, + bool use_mmap) + : segment_id_(segment_id), + key_(fmt::format("seg_{}_f_{}", segment_id, field_data_info.field_id)), + use_mmap_(use_mmap), + field_meta_(field_meta), + meta_(use_mmap ? milvus::cachinglayer::StorageType::DISK + : milvus::cachinglayer::StorageType::MEMORY) { + meta_.num_rows_until_chunk_.push_back(0); + meta_.num_rows_until_chunk_.push_back(field_data_info.row_count); +} + +DefaultValueChunkTranslator::~DefaultValueChunkTranslator() { +} + +size_t +DefaultValueChunkTranslator::num_cells() const { + return 1; +} + +milvus::cachinglayer::cid_t +DefaultValueChunkTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const { + return 0; +} + +milvus::cachinglayer::ResourceUsage +DefaultValueChunkTranslator::estimated_byte_size_of_cell( + milvus::cachinglayer::cid_t cid) const { + return milvus::cachinglayer::ResourceUsage{0, 0}; +} + +const std::string& +DefaultValueChunkTranslator::key() const { + return key_; +} + +std::vector< + std::pair>> +DefaultValueChunkTranslator::get_cells( + const std::vector& cids) { + AssertInfo(cids.size() == 1 && cids[0] == 0, + "DefaultValueChunkTranslator only supports one cell"); + auto num_rows = meta_.num_rows_until_chunk_[1]; + auto builder = + milvus::storage::CreateArrowBuilder(field_meta_.get_data_type()); + arrow::Status ast; + if (field_meta_.default_value().has_value()) { + builder->Reserve(num_rows); + auto scalar = storage::CreateArrowScalarFromDefaultValue(field_meta_); + ast = builder->AppendScalar(*scalar, num_rows); + } else { + ast = builder->AppendNulls(num_rows); + } + AssertInfo(ast.ok(), + "append null/default values to arrow builder failed: {}", + ast.ToString()); + arrow::ArrayVector array_vec; + array_vec.emplace_back(builder->Finish().ValueOrDie()); + auto chunk = create_chunk( + field_meta_, + IsVectorDataType(field_meta_.get_data_type()) && + !IsSparseFloatVectorDataType(field_meta_.get_data_type()) + ? field_meta_.get_dim() + : 1, + array_vec); + + std::vector< + std::pair>> + res; + res.reserve(1); + res.emplace_back(0, std::move(chunk)); + return res; +} + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.h b/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.h new file mode 100644 index 0000000000..a02327e27a --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.h @@ -0,0 +1,62 @@ +// Copyright (C) 2019-2025 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" +#include "common/Chunk.h" +#include "common/FieldMeta.h" +#include "mmap/Types.h" +#include "segcore/storagev1translator/ChunkTranslator.h" + +namespace milvus::segcore::storagev1translator { + +// This is to support add field. +class DefaultValueChunkTranslator + : public milvus::cachinglayer::Translator { + public: + DefaultValueChunkTranslator(int64_t segment_id, + FieldMeta field_meta, + FieldDataInfo field_data_info, + bool use_mmap); + + ~DefaultValueChunkTranslator() override; + + size_t + num_cells() const override; + milvus::cachinglayer::cid_t + cell_id_of(milvus::cachinglayer::uid_t uid) const override; + milvus::cachinglayer::ResourceUsage + estimated_byte_size_of_cell(milvus::cachinglayer::cid_t cid) const override; + const std::string& + key() const override; + std::vector< + std::pair>> + get_cells(const std::vector& cids) override; + + milvus::cachinglayer::Meta* + meta() override { + return &meta_; + } + + private: + int64_t segment_id_; + std::string key_; + bool use_mmap_; + CTMeta meta_; + milvus::FieldMeta field_meta_; +}; + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.cpp b/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.cpp new file mode 100644 index 0000000000..146f7caeff --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.cpp @@ -0,0 +1,139 @@ +#include "segcore/storagev1translator/InsertRecordTranslator.h" + +#include +#include +#include + +#include "fmt/core.h" + +#include "cachinglayer/Utils.h" +#include "common/ChunkWriter.h" +#include "common/Types.h" +#include "common/SystemProperty.h" +#include "segcore/Utils.h" +#include "storage/ThreadPools.h" + +namespace milvus::segcore::storagev1translator { + +InsertRecordTranslator::InsertRecordTranslator( + int64_t segment_id, + DataType data_type, + FieldDataInfo field_data_info, + SchemaPtr schema, + bool is_sorted_by_pk, + std::vector insert_files, + ChunkedSegmentSealedImpl* chunked_segment) + : segment_id_(segment_id), + data_type_(data_type), + key_(fmt::format("seg_{}_ir_f_{}", segment_id, field_data_info.field_id)), + field_data_info_(field_data_info), + schema_(schema), + is_sorted_by_pk_(is_sorted_by_pk), + insert_files_(insert_files), + chunked_segment_(chunked_segment), + meta_(milvus::cachinglayer::StorageType::MEMORY) { +} + +size_t +InsertRecordTranslator::num_cells() const { + return 1; +} + +milvus::cachinglayer::cid_t +InsertRecordTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const { + return 0; +} + +milvus::cachinglayer::ResourceUsage +InsertRecordTranslator::estimated_byte_size_of_cell( + milvus::cachinglayer::cid_t cid) const { + return {0, 0}; +} + +const std::string& +InsertRecordTranslator::key() const { + return key_; +} + +std::vector>>> +InsertRecordTranslator::get_cells( + const std::vector& cids) { + AssertInfo(cids.size() == 1 && cids[0] == 0, + "InsertRecordTranslator only supports single cell"); + FieldId fid = FieldId(field_data_info_.field_id); + auto parallel_degree = + static_cast(DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE); + // TODO(tiered storage 4): we should phase out this thread pool and use folly executor. + auto& pool = ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE); + pool.Submit(LoadArrowReaderFromRemote, + insert_files_, + field_data_info_.arrow_reader_channel); + LOG_INFO("segment {} submits load field {} task to thread pool", + segment_id_, + field_data_info_.field_id); + auto num_rows = field_data_info_.row_count; + AssertInfo(milvus::SystemProperty::Instance().IsSystem(fid), + "system field is not system field"); + auto system_field_type = + milvus::SystemProperty::Instance().GetSystemFieldType(fid); + AssertInfo(system_field_type == SystemFieldType::Timestamp, + "system field is not timestamp"); + std::vector timestamps(num_rows); + int64_t offset = 0; + FieldMeta field_meta( + FieldName(""), FieldId(0), DataType::INT64, false, std::nullopt); + + std::shared_ptr r; + while (field_data_info_.arrow_reader_channel->pop(r)) { + arrow::ArrayVector array_vec = read_single_column_batches(r->reader); + auto chunk = create_chunk(field_meta, 1, array_vec); + auto chunk_ptr = static_cast(chunk.get()); + std::copy_n(static_cast(chunk_ptr->Span().data()), + chunk_ptr->Span().row_count(), + timestamps.data() + offset); + offset += chunk_ptr->Span().row_count(); + } + + TimestampIndex index; + auto min_slice_length = num_rows < 4096 ? 1 : 4096; + auto meta = + GenerateFakeSlices(timestamps.data(), num_rows, min_slice_length); + index.set_length_meta(std::move(meta)); + // todo ::opt to avoid copy timestamps from field data + index.build_with(timestamps.data(), num_rows); + + std::unique_ptr> ir = + std::make_unique>(*schema_, + MAX_ROW_COUNT); + + // use special index + AssertInfo(ir->timestamps_.empty(), "already exists"); + ir->timestamps_.set_data_raw(0, timestamps.data(), timestamps.size()); + ir->timestamp_index_ = std::move(index); + AssertInfo(ir->timestamps_.num_chunk() == 1, + "num chunk not equal to 1 for sealed segment"); + chunked_segment_->stats_.mem_size += sizeof(Timestamp) * num_rows; + + auto pk_field_id = schema_->get_primary_field_id(); + AssertInfo(pk_field_id.has_value(), + "primary key field not found in schema"); + auto pk_field_meta = schema_->operator[](pk_field_id.value()); + + // set pks to offset + if (!is_sorted_by_pk_) { + AssertInfo(ir->empty_pks(), "already exists"); + auto it = chunked_segment_->fields_.find(pk_field_id.value()); + AssertInfo(it != chunked_segment_->fields_.end(), + "primary key field not found in segment"); + ir->insert_pks(pk_field_meta.get_data_type(), it->second.get()); + ir->seal_pks(); + } + std::vector>>> + cells; + cells.emplace_back(0, std::move(ir)); + return cells; +} + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.h b/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.h new file mode 100644 index 0000000000..9a41147de3 --- /dev/null +++ b/internal/core/src/segcore/storagev1translator/InsertRecordTranslator.h @@ -0,0 +1,72 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" +#include "common/Schema.h" +#include "segcore/InsertRecord.h" +#include "segcore/ChunkedSegmentSealedImpl.h" + +namespace milvus::segcore::storagev1translator { + +class InsertRecordTranslator : public milvus::cachinglayer::Translator< + milvus::segcore::InsertRecord> { + public: + InsertRecordTranslator(int64_t segment_id, + DataType data_type, + FieldDataInfo field_data_info, + SchemaPtr schema, + bool is_sorted_by_pk, + std::vector insert_files, + ChunkedSegmentSealedImpl* chunked_segment); + + size_t + num_cells() const override; + milvus::cachinglayer::cid_t + cell_id_of(milvus::cachinglayer::uid_t uid) const override; + milvus::cachinglayer::ResourceUsage + estimated_byte_size_of_cell(milvus::cachinglayer::cid_t cid) const override; + const std::string& + key() const override; + + // each calling of this will trigger a new download. + std::vector>>> + get_cells(const std::vector& cids) override; + + // InsertRecord does not have meta. + milvus::cachinglayer::Meta* + meta() override { + return &meta_; + } + + private: + std::unique_ptr> + load_column_in_memory() const; + + int64_t segment_id_; + std::string key_; + DataType data_type_; + FieldDataInfo field_data_info_; + std::vector insert_files_; + mutable size_t estimated_byte_size_of_cell_; + SchemaPtr schema_; + bool is_sorted_by_pk_; + ChunkedSegmentSealedImpl* chunked_segment_; + milvus::cachinglayer::Meta meta_; +}; + +} // namespace milvus::segcore::storagev1translator diff --git a/internal/core/src/storage/ChunkCache.cpp b/internal/core/src/storage/ChunkCache.cpp deleted file mode 100644 index 5442a05a94..0000000000 --- a/internal/core/src/storage/ChunkCache.cpp +++ /dev/null @@ -1,202 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include -#include -#include - -#include "ChunkCache.h" -#include "boost/filesystem/operations.hpp" -#include "boost/filesystem/path.hpp" -#include "common/Chunk.h" -#include "common/ChunkWriter.h" -#include "common/FieldMeta.h" -#include "common/Types.h" -#include "log/Log.h" - -namespace milvus::storage { -std::shared_ptr -ChunkCache::Read(const std::string& filepath, - const FieldMeta& field_meta, - bool mmap_enabled, - bool mmap_rss_not_need) { - // use rlock to get future - { - std::shared_lock lck(mutex_); - auto it = columns_.find(filepath); - if (it != columns_.end()) { - lck.unlock(); - auto result = it->second.second.get(); - AssertInfo(result, "unexpected null column, file={}", filepath); - return result; - } - } - - // lock for mutation - std::unique_lock lck(mutex_); - // double check no-futurn - auto it = columns_.find(filepath); - if (it != columns_.end()) { - lck.unlock(); - auto result = it->second.second.get(); - AssertInfo(result, "unexpected null column, file={}", filepath); - return result; - } - - std::promise> p; - std::shared_future> f = p.get_future(); - columns_.emplace(filepath, std::make_pair(std::move(p), f)); - lck.unlock(); - - // release lock and perform download and decode - // other thread request same path shall get the future. - bool allocate_success = false; - ErrorCode err_code = Success; - std::string err_msg = ""; - std::shared_ptr column; - try { - auto field_data = - DownloadAndDecodeRemoteFile(cm_.get(), filepath, false); - - std::shared_ptr chunk; - auto dim = IsSparseFloatVectorDataType(field_meta.get_data_type()) - ? 1 - : field_meta.get_dim(); - if (mmap_enabled) { - auto path = std::filesystem::path(CachePath(filepath)); - auto dir = path.parent_path(); - std::filesystem::create_directories(dir); - - auto file = File::Open(path.string(), O_CREAT | O_TRUNC | O_RDWR); - - arrow::ArrayVector array_vec = - read_single_column_batches(field_data->GetReader()->reader); - chunk = create_chunk(field_meta, dim, file, 0, array_vec); - // unlink - auto ok = unlink(path.c_str()); - AssertInfo(ok == 0, - "failed to unlink mmap data file {}, err: {}", - path.c_str(), - strerror(errno)); - } else { - arrow::ArrayVector array_vec = - read_single_column_batches(field_data->GetReader()->reader); - chunk = create_chunk(field_meta, dim, array_vec); - } - - auto data_type = field_meta.get_data_type(); - if (IsSparseFloatVectorDataType(data_type)) { - auto sparse_column = - std::make_shared(field_meta); - sparse_column->AddChunk(chunk); - column = std::move(sparse_column); - } else if (IsVariableDataType(data_type)) { - AssertInfo(false, - "TODO: unimplemented for variable data type: {}", - data_type); - } else { - std::vector> chunks{chunk}; - column = std::make_shared(field_meta, chunks); - } - if (mmap_enabled && mmap_rss_not_need) { - auto ok = madvise(reinterpret_cast( - const_cast(column->MmappedData())), - column->DataByteSize(), - ReadAheadPolicy_Map["dontneed"]); - if (ok != 0) { - LOG_WARN( - "failed to madvise to the data file {}, addr {}, size {}, " - "err: " - "{}", - filepath, - static_cast(column->MmappedData()), - column->DataByteSize(), - strerror(errno)); - } - } - } catch (const SegcoreError& e) { - err_code = e.get_error_code(); - err_msg = fmt::format("failed to read for chunkCache, seg_core_err:{}", - e.what()); - } - std::unique_lock mmap_lck(mutex_); - - it = columns_.find(filepath); - if (it != columns_.end()) { - // check pair exists then set value - it->second.first.set_value(column); - if (allocate_success) { - AssertInfo(column, "unexpected null column, file={}", filepath); - } - } else { - PanicInfo(UnexpectedError, - "Wrong code, the thread to download for cache should get the " - "target entry"); - } - if (err_code != Success) { - columns_.erase(filepath); - throw SegcoreError(err_code, err_msg); - } - return column; -} - -void -ChunkCache::Remove(const std::string& filepath) { - std::unique_lock lck(mutex_); - columns_.erase(filepath); -} - -void -ChunkCache::Prefetch(const std::string& filepath) { - std::shared_lock lck(mutex_); - auto it = columns_.find(filepath); - if (it == columns_.end()) { - return; - } - - auto column = it->second.second.get(); - auto ok = madvise( - reinterpret_cast(const_cast(column->MmappedData())), - column->DataByteSize(), - read_ahead_policy_); - if (ok != 0) { - LOG_WARN( - "failed to madvise to the data file {}, addr {}, size {}, err: {}", - filepath, - static_cast(column->MmappedData()), - column->DataByteSize(), - strerror(errno)); - } -} - -// TODO(sunby): use mmap chunk manager to create chunk -std::string -ChunkCache::CachePath(const std::string& filepath) { - auto path = std::filesystem::path(filepath); - auto prefix = std::filesystem::path(path_prefix_); - - // Cache path shall not use absolute filepath direct, it shall always under path_prefix_ - if (path.is_absolute()) { - return (prefix / - filepath.substr(path.root_directory().string().length(), - filepath.length())) - .string(); - } - - return (prefix / filepath).string(); -} - -} // namespace milvus::storage diff --git a/internal/core/src/storage/ChunkCache.h b/internal/core/src/storage/ChunkCache.h deleted file mode 100644 index 20584b4159..0000000000 --- a/internal/core/src/storage/ChunkCache.h +++ /dev/null @@ -1,83 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#pragma once -#include -#include -#include "common/FieldMeta.h" -#include "storage/MmapChunkManager.h" -#include "mmap/ChunkedColumn.h" - -namespace milvus::storage { - -extern std::map ReadAheadPolicy_Map; - -class ChunkCache { - public: - explicit ChunkCache(std::string& path_prefix, - const std::string& read_ahead_policy, - ChunkManagerPtr cm, - MmapChunkManagerPtr mcm) - : path_prefix_(path_prefix), cm_(cm), mcm_(mcm) { - auto iter = ReadAheadPolicy_Map.find(read_ahead_policy); - AssertInfo(iter != ReadAheadPolicy_Map.end(), - "unrecognized read ahead policy: {}, " - "should be one of `normal, random, sequential, " - "willneed, dontneed`", - read_ahead_policy); - read_ahead_policy_ = iter->second; - LOG_INFO("Init ChunkCache with read_ahead_policy: {}", - read_ahead_policy); - } - - ~ChunkCache() = default; - - public: - std::shared_ptr - Read(const std::string& filepath, - const FieldMeta& field_meta, - bool mmap_enabled, - bool mmap_rss_not_need = false); - - void - Remove(const std::string& filepath); - - void - Prefetch(const std::string& filepath); - - private: - std::string - CachePath(const std::string& filepath); - - private: - using ColumnTable = std::unordered_map< - std::string, - std::pair>, - std::shared_future>>>; - - private: - mutable std::shared_mutex mutex_; - int read_ahead_policy_; - ChunkManagerPtr cm_; - MmapChunkManagerPtr mcm_; - ColumnTable columns_; - - std::string path_prefix_; -}; - -using ChunkCachePtr = std::shared_ptr; - -} // namespace milvus::storage diff --git a/internal/core/src/storage/Event.cpp b/internal/core/src/storage/Event.cpp index beb1855192..f839e52d34 100644 --- a/internal/core/src/storage/Event.cpp +++ b/internal/core/src/storage/Event.cpp @@ -359,7 +359,7 @@ BaseEvent::Serialize() { int header_size = header.size(); int len = header_size + data_size; - std::vector res(len); + std::vector res(len, 0); int offset = 0; memcpy(res.data() + offset, header.data(), header_size); offset += header_size; @@ -384,7 +384,7 @@ DescriptorEvent::Serialize() { int header_size = header.size(); int len = header_size + data_size + sizeof(MAGIC_NUM); - std::vector res(len); + std::vector res(len, 0); int offset = 0; memcpy(res.data(), &MAGIC_NUM, sizeof(MAGIC_NUM)); offset += sizeof(MAGIC_NUM); diff --git a/internal/core/src/storage/MmapManager.h b/internal/core/src/storage/MmapManager.h index 8043474aed..2c79dc3707 100644 --- a/internal/core/src/storage/MmapManager.h +++ b/internal/core/src/storage/MmapManager.h @@ -16,9 +16,8 @@ #pragma once -#include -#include "ChunkCache.h" -#include "RemoteChunkManagerSingleton.h" +#include "storage/MmapChunkManager.h" +#include "storage/Types.h" namespace milvus::storage { /** @@ -42,9 +41,6 @@ class MmapManager { return instance; } ~MmapManager() { - if (cc_ != nullptr) { - cc_ = nullptr; - } // delete mmap chunk manager at last if (mcm_ != nullptr) { mcm_ = nullptr; @@ -62,15 +58,6 @@ class MmapManager { mmap_config_.disk_limit, mmap_config_.fix_file_size); } - if (cc_ == nullptr) { - auto rcm = RemoteChunkManagerSingleton::GetInstance() - .GetRemoteChunkManager(); - cc_ = std::make_shared( - mmap_config_.mmap_path, - std::move(mmap_config_.cache_read_ahead_policy), - rcm, - mcm_); - } LOG_INFO("Init MmapConfig with MmapConfig: {}", mmap_config_.ToString()); init_flag_ = true; @@ -79,12 +66,6 @@ class MmapManager { } } - ChunkCachePtr - GetChunkCache() { - AssertInfo(init_flag_ == true, "Mmap manager has not been init."); - return cc_; - } - MmapChunkManagerPtr GetMmapChunkManager() { AssertInfo(init_flag_ == true, "Mmap manager has not been init."); @@ -119,7 +100,6 @@ class MmapManager { mutable std::mutex init_mutex_; MmapConfig mmap_config_; MmapChunkManagerPtr mcm_ = nullptr; - ChunkCachePtr cc_ = nullptr; std::atomic init_flag_ = false; }; diff --git a/internal/core/src/storage/ThreadPool.h b/internal/core/src/storage/ThreadPool.h index 521ddd9c89..74a1990a0e 100644 --- a/internal/core/src/storage/ThreadPool.h +++ b/internal/core/src/storage/ThreadPool.h @@ -85,7 +85,6 @@ class ThreadPool { template auto - // Submit(F&& f, Args&&... args) -> std::future; Submit(F&& f, Args&&... args) -> std::future { std::function func = std::bind(std::forward(f), std::forward(args)...); diff --git a/internal/core/src/storage/Util.cpp b/internal/core/src/storage/Util.cpp index 3126880511..8d9dbb16fb 100644 --- a/internal/core/src/storage/Util.cpp +++ b/internal/core/src/storage/Util.cpp @@ -18,6 +18,7 @@ #include #include "arrow/array/builder_binary.h" +#include "arrow/scalar.h" #include "arrow/type_fwd.h" #include "fmt/format.h" #include "log/Log.h" @@ -326,6 +327,42 @@ CreateArrowBuilder(DataType data_type, int dim) { } } +std::shared_ptr +CreateArrowScalarFromDefaultValue(const FieldMeta& field_meta) { + auto default_var = field_meta.default_value(); + AssertInfo(default_var.has_value(), + "cannot create Arrow Scalar from empty default value"); + auto default_value = default_var.value(); + switch (field_meta.get_data_type()) { + case DataType::BOOL: + return std::make_shared( + default_value.bool_data()); + case DataType::INT8: + case DataType::INT16: + case DataType::INT32: + return std::make_shared( + default_value.int_data()); + case DataType::INT64: + return std::make_shared( + default_value.long_data()); + case DataType::FLOAT: + return std::make_shared( + default_value.float_data()); + case DataType::DOUBLE: + return std::make_shared( + default_value.double_data()); + case DataType::VARCHAR: + case DataType::STRING: + case DataType::TEXT: + return std::make_shared( + default_value.string_data()); + default: + PanicInfo(DataTypeInvalid, + "unsupported default value data type {}", + field_meta.get_data_type()); + } +} + std::shared_ptr CreateArrowSchema(DataType data_type, bool nullable) { switch (static_cast(data_type)) { @@ -649,31 +686,6 @@ EncodeAndUploadIndexSlice(ChunkManager* chunk_manager, return std::make_pair(std::move(object_key), serialized_index_size); } -std::pair -EncodeAndUploadFieldSlice(ChunkManager* chunk_manager, - void* buf, - int64_t element_count, - FieldDataMeta field_data_meta, - const FieldMeta& field_meta, - std::string object_key) { - // dim should not be used for sparse float vector field - auto dim = IsSparseFloatVectorDataType(field_meta.get_data_type()) - ? -1 - : field_meta.get_dim(); - auto field_data = - CreateFieldData(field_meta.get_data_type(), false, dim, 0); - field_data->FillFieldData(buf, element_count); - auto payload_reader = std::make_shared(field_data); - auto insertData = std::make_shared(payload_reader); - insertData->SetFieldDataMeta(field_data_meta); - auto serialized_inserted_data = insertData->serialize_to_remote_file(); - auto serialized_inserted_data_size = serialized_inserted_data.size(); - chunk_manager->Write(object_key, - serialized_inserted_data.data(), - serialized_inserted_data_size); - return std::make_pair(std::move(object_key), serialized_inserted_data_size); -} - std::vector>> GetObjectData(ChunkManager* remote_chunk_manager, const std::vector& remote_files) { diff --git a/internal/core/src/storage/Util.h b/internal/core/src/storage/Util.h index 09c7922638..04ca7bb4f8 100644 --- a/internal/core/src/storage/Util.h +++ b/internal/core/src/storage/Util.h @@ -57,6 +57,19 @@ CreateArrowBuilder(DataType data_type); std::shared_ptr CreateArrowBuilder(DataType data_type, int dim); +/// \brief Utility function to create arrow:Scalar from FieldMeta.default_value +/// +/// Construct a arrow::Scalar based on input field meta +/// The data_type_ is checked to determine which `one_of` member of default value shall be used +/// Note that: +/// 1. default_value shall have value +/// 2. the type check shall be guaranteed(current by go side) +/// +/// \param[in] field_meta the field meta object to construct arrow::Scalar from. +/// \return an std::shared_ptr of arrow::Scalar +std::shared_ptr +CreateArrowScalarFromDefaultValue(const FieldMeta& field_meta); + std::shared_ptr CreateArrowSchema(DataType data_type, bool nullable); @@ -131,14 +144,6 @@ EncodeAndUploadIndexSlice(ChunkManager* chunk_manager, FieldDataMeta field_meta, std::string object_key); -std::pair -EncodeAndUploadFieldSlice(ChunkManager* chunk_manager, - void* buf, - int64_t element_count, - FieldDataMeta field_data_meta, - const FieldMeta& field_meta, - std::string object_key); - std::vector>> GetObjectData(ChunkManager* remote_chunk_manager, const std::vector& remote_files); @@ -160,9 +165,6 @@ GetNumRowsForLoadInfo(const LoadFieldDataInfo& load_info); void ReleaseArrowUnused(); -// size_t -// getCurrentRSS(); - ChunkManagerPtr CreateChunkManager(const StorageConfig& storage_config); @@ -208,6 +210,7 @@ SortByPath(std::vector& paths) { }); } +// used only for test inline std::shared_ptr ConvertFieldDataToArrowDataWrapper(const FieldDataPtr& field_data) { BaseEventData event_data; diff --git a/internal/core/unittest/CMakeLists.txt b/internal/core/unittest/CMakeLists.txt index da95a1226b..55b394fc9d 100644 --- a/internal/core/unittest/CMakeLists.txt +++ b/internal/core/unittest/CMakeLists.txt @@ -37,7 +37,6 @@ set(MILVUS_TEST_FILES test_bitmap_index.cpp test_bool_index.cpp test_c_api.cpp - test_chunk_cache.cpp test_chunk.cpp test_chunk_vector.cpp test_common.cpp @@ -121,7 +120,8 @@ if (LINUX OR APPLE) ${MILVUS_TEST_FILES} test_scalar_index_creator.cpp test_string_index.cpp - test_array.cpp test_array_expr.cpp) + test_array.cpp + test_array_expr.cpp) endif() if (DEFINED AZURE_BUILD_DIR) @@ -177,6 +177,7 @@ add_executable(all_tests target_link_libraries(all_tests gtest + gmock milvus_core knowhere milvus-storage @@ -185,26 +186,27 @@ target_link_libraries(all_tests install(TARGETS all_tests DESTINATION unittest) add_subdirectory(bench) +add_subdirectory(test_cachinglayer) # if (USE_DYNAMIC_SIMD) # add_executable(dynamic_simd_test # test_simd.cpp) -# +# # target_link_libraries(dynamic_simd_test # milvus_simd # milvus_log # gtest # ${CONAN_LIBS}) -# +# # install(TARGETS dynamic_simd_test DESTINATION unittest) # endif() -add_executable(bitset_test +add_executable(bitset_test test_bitset.cpp ) -target_link_libraries(bitset_test - milvus_bitset - gtest +target_link_libraries(bitset_test + milvus_bitset + gtest ${CONAN_LIBS} ) install(TARGETS bitset_test DESTINATION unittest) diff --git a/internal/core/unittest/bench/bench_search.cpp b/internal/core/unittest/bench/bench_search.cpp index 0049752904..f93d7def89 100644 --- a/internal/core/unittest/bench/bench_search.cpp +++ b/internal/core/unittest/bench/bench_search.cpp @@ -12,9 +12,12 @@ #include #include #include +#include "common/type_c.h" +#include "segcore/segment_c.h" #include "segcore/SegmentGrowing.h" #include "segcore/SegmentSealed.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::query; @@ -109,7 +112,12 @@ Search_Sealed(benchmark::State& state) { auto dataset_ = DataGen(schema, N); return dataset_; }(); - SealedLoadFieldData(dataset_, *segment); + auto storage_config = get_default_local_storage_config(); + auto cm = storage::CreateChunkManager(storage_config); + auto load_info = PrepareInsertBinlog(1, 1, 1, dataset_, cm); + auto segment_ptr = segment.get(); + auto status = LoadFieldData(segment_ptr, &load_info); + ASSERT_EQ(status.error_code, Success); auto choice = state.range(0); if (choice == 0) { // Brute Force diff --git a/internal/core/unittest/init_gtest.cpp b/internal/core/unittest/init_gtest.cpp index adc1b3b683..ff57bf0ece 100644 --- a/internal/core/unittest/init_gtest.cpp +++ b/internal/core/unittest/init_gtest.cpp @@ -28,5 +28,8 @@ main(int argc, char** argv) { get_default_local_storage_config()); milvus::storage::MmapManager::GetInstance().Init(get_default_mmap_config()); + milvus::cachinglayer::Manager::ConfigureTieredStorage( + true, 1024 * 1024 * 1024, 1024 * 1024 * 1024); + return RUN_ALL_TESTS(); } diff --git a/internal/core/unittest/test_array_inverted_index.cpp b/internal/core/unittest/test_array_inverted_index.cpp index 42958c5f08..1da3d87148 100644 --- a/internal/core/unittest/test_array_inverted_index.cpp +++ b/internal/core/unittest/test_array_inverted_index.cpp @@ -21,6 +21,8 @@ #include "query/PlanProto.h" #include "query/ExecPlanNodeVisitor.h" +#include "test_utils/storage_test_utils.h" + using namespace milvus; using namespace milvus::query; using namespace milvus::segcore; @@ -61,7 +63,6 @@ class ArrayInvertedIndexTest : public ::testing::Test { void SetUp() override { schema_ = GenTestSchema(); - seg_ = CreateSealedSegment(schema_); N_ = 3000; uint64_t seed = 19190504; auto raw_data = DataGen(schema_, N_, seed); @@ -100,7 +101,7 @@ class ArrayInvertedIndexTest : public ::testing::Test { } vec_of_array_.push_back(array); } - SealedLoadFieldData(raw_data, *seg_); + seg_ = CreateSealedWithFieldDataLoaded(schema_, raw_data); LoadInvertedIndex(); } diff --git a/internal/core/unittest/test_binlog_index.cpp b/internal/core/unittest/test_binlog_index.cpp index e1eea0a891..9fe4f53507 100644 --- a/internal/core/unittest/test_binlog_index.cpp +++ b/internal/core/unittest/test_binlog_index.cpp @@ -12,21 +12,18 @@ #include #include #include -#include #include "index/IndexFactory.h" -#include "knowhere/comp/brute_force.h" #include "pb/plan.pb.h" -#include "pb/schema.pb.h" #include "query/Plan.h" #include "segcore/segcore_init_c.h" #include "segcore/SegmentSealed.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::segcore; -namespace pb = milvus::proto; std::unique_ptr GenRandomFloatVecData(int rows, int dim, int seed = 42) { @@ -121,41 +118,30 @@ class BinlogIndexTest : public ::testing::TestWithParam { void LoadOtherFields() { auto dataset = DataGen(schema, data_n); - // load id - LoadFieldDataInfo row_id_info; - FieldMeta row_id_field_meta(FieldName("RowID"), - RowFieldID, - DataType::INT64, - false, - std::nullopt); - auto field_data = std::make_shared>( - DataType::INT64, false); - field_data->FillFieldData(dataset.row_ids_.data(), data_n); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo{ - RowFieldID.get(), - data_n, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(RowFieldID, field_data_info); - // load ts - LoadFieldDataInfo ts_info; - FieldMeta ts_field_meta(FieldName("Timestamp"), - TimestampFieldID, - DataType::INT64, - false, - std::nullopt); - field_data = std::make_shared>( - DataType::INT64, false); - field_data->FillFieldData(dataset.timestamps_.data(), data_n); - auto arrow_data_wrapper2 = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - field_data_info = - FieldDataInfo{TimestampFieldID.get(), - data_n, - std::vector>{ - arrow_data_wrapper2}}; - segment->LoadFieldData(TimestampFieldID, field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + dataset, + cm, + "", + {vec_field_id.get()}); + segment->LoadFieldData(load_info); + } + + void + LoadVectorField(std::string mmap_dir_path = "") { + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + vec_field_id.get(), + {vec_field_data}, + cm, + mmap_dir_path); + segment->LoadFieldData(load_info); } protected: @@ -196,13 +182,7 @@ TEST_P(BinlogIndexTest, AccuracyWithLoadFieldData) { segcore_config.set_enable_interim_segment_index(true); segcore_config.set_nprobe(32); // 1. load field data, and build binlog index for binlog data - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(vec_field_data); - auto field_data_info = FieldDataInfo{ - vec_field_id.get(), - data_n, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(vec_field_id, field_data_info); + LoadVectorField(); //assert segment has been built binlog index EXPECT_TRUE(segment->HasIndex(vec_field_id)); @@ -294,15 +274,7 @@ TEST_P(BinlogIndexTest, AccuracyWithMapFieldData) { segcore_config.set_enable_interim_segment_index(true); segcore_config.set_nprobe(32); // 1. load field data, and build binlog index for binlog data - FieldDataInfo field_data_info; - field_data_info.field_id = vec_field_id.get(); - field_data_info.row_count = data_n; - field_data_info.mmap_dir_path = "./data/mmap-test"; - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(vec_field_data); - field_data_info.arrow_reader_channel->push(arrow_data_wrapper); - field_data_info.arrow_reader_channel->close(); - segment->MapFieldData(vec_field_id, field_data_info); + LoadVectorField("./data/mmap-test"); //assert segment has been built binlog index EXPECT_TRUE(segment->HasIndex(vec_field_id)); @@ -392,13 +364,7 @@ TEST_P(BinlogIndexTest, DisableInterimIndex) { LoadOtherFields(); SegcoreSetEnableTempSegmentIndex(false); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(vec_field_data); - auto field_data_info = FieldDataInfo{ - vec_field_id.get(), - data_n, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(vec_field_id, field_data_info); + LoadVectorField(); EXPECT_FALSE(segment->HasIndex(vec_field_id)); EXPECT_EQ(segment->get_row_count(), data_n); @@ -439,15 +405,7 @@ TEST_P(BinlogIndexTest, LoadBingLogWihIDMAP) { segment = CreateSealedSegment(schema, collection_index_meta); LoadOtherFields(); - - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(vec_field_data); - - auto field_data_info = FieldDataInfo{ - vec_field_id.get(), - data_n, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(vec_field_id, field_data_info); + LoadVectorField(); EXPECT_FALSE(segment->HasIndex(vec_field_id)); EXPECT_EQ(segment->get_row_count(), data_n); @@ -461,13 +419,7 @@ TEST_P(BinlogIndexTest, LoadBinlogWithoutIndexMeta) { segment = CreateSealedSegment(schema, collection_index_meta); SegcoreSetEnableTempSegmentIndex(true); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(vec_field_data); - auto field_data_info = FieldDataInfo{ - vec_field_id.get(), - data_n, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(vec_field_id, field_data_info); + LoadVectorField(); EXPECT_FALSE(segment->HasIndex(vec_field_id)); EXPECT_EQ(segment->get_row_count(), data_n); diff --git a/internal/core/unittest/test_c_api.cpp b/internal/core/unittest/test_c_api.cpp index d7f2b46e1f..00b6bc4f2d 100644 --- a/internal/core/unittest/test_c_api.cpp +++ b/internal/core/unittest/test_c_api.cpp @@ -47,11 +47,10 @@ #include "plan/PlanNode.h" #include "segcore/load_index_c.h" #include "test_utils/c_api_test_utils.h" +#include "test_utils/DataGen.h" #include "segcore/vector_index_c.h" #include "common/jsmn.h" -namespace chrono = std::chrono; - using namespace milvus; using namespace milvus::test; using namespace milvus::index; @@ -437,13 +436,8 @@ TEST(CApiTest, DeleteTest) { auto delete_data = serialize(ids.get()); uint64_t delete_timestamps[] = {0, 0, 0}; - auto offset = 0; - auto del_res = Delete(segment, - offset, - 3, - delete_data.data(), - delete_data.size(), - delete_timestamps); + auto del_res = Delete( + segment, 3, delete_data.data(), delete_data.size(), delete_timestamps); ASSERT_EQ(del_res.error_code, Success); DeleteCollection(collection); @@ -480,9 +474,7 @@ TEST(CApiTest, MultiDeleteGrowingSegment) { delete_pks.end()); auto delete_data = serialize(ids.get()); std::vector delete_timestamps(1, dataset.timestamps_[N - 1]); - offset = 0; auto del_res = Delete(segment, - offset, 1, delete_data.data(), delete_data.size(), @@ -546,9 +538,7 @@ TEST(CApiTest, MultiDeleteGrowingSegment) { delete_pks.end()); delete_data = serialize(ids.get()); delete_timestamps[0]++; - offset = 0; del_res = Delete(segment, - offset, 1, delete_data.data(), delete_data.size(), @@ -598,7 +588,11 @@ TEST(CApiTest, MultiDeleteSealedSegment) { auto segment_interface = reinterpret_cast(segment); auto sealed_segment = dynamic_cast(segment_interface); - SealedLoadFieldData(dataset, *sealed_segment); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareInsertBinlog( + kCollectionID, kPartitionID, kSegmentID, dataset, cm); + sealed_segment->LoadFieldData(load_info); // delete data pks = {1} std::vector delete_pks = {1}; @@ -607,9 +601,7 @@ TEST(CApiTest, MultiDeleteSealedSegment) { delete_pks.end()); auto delete_data = serialize(ids.get()); std::vector delete_timestamps(1, dataset.timestamps_[N - 1]); - auto offset = 0; auto del_res = Delete(segment, - offset, 1, delete_data.data(), delete_data.size(), @@ -674,9 +666,7 @@ TEST(CApiTest, MultiDeleteSealedSegment) { delete_pks.end()); delete_data = serialize(ids.get()); delete_timestamps[0]++; - offset = 0; del_res = Delete(segment, - offset, 1, delete_data.data(), delete_data.size(), @@ -774,9 +764,7 @@ TEST(CApiTest, DeleteRepeatedPksFromGrowingSegment) { auto delete_data = serialize(ids.get()); std::vector delete_timestamps(3, dataset.timestamps_[N - 1]); - offset = 0; auto del_res = Delete(segment, - offset, 3, delete_data.data(), delete_data.size(), @@ -814,7 +802,11 @@ TEST(CApiTest, DeleteRepeatedPksFromSealedSegment) { auto segment_interface = reinterpret_cast(segment); auto sealed_segment = dynamic_cast(segment_interface); - SealedLoadFieldData(dataset, *sealed_segment); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareInsertBinlog( + kCollectionID, kPartitionID, kSegmentID, dataset, cm); + sealed_segment->LoadFieldData(load_info); std::vector retrive_row_ids; // create retrieve plan pks in {1, 2, 3} @@ -856,10 +848,7 @@ TEST(CApiTest, DeleteRepeatedPksFromSealedSegment) { auto delete_data = serialize(ids.get()); std::vector delete_timestamps(3, dataset.timestamps_[N - 1]); - auto offset = 0; - auto del_res = Delete(segment, - offset, 3, delete_data.data(), delete_data.size(), @@ -984,10 +973,7 @@ TEST(CApiTest, InsertSamePkAfterDeleteOnGrowingSegment) { auto delete_data = serialize(ids.get()); std::vector delete_timestamps(3, dataset.timestamps_[N - 1]); - offset = 0; - auto del_res = Delete(segment, - offset, 3, delete_data.data(), delete_data.size(), @@ -1071,8 +1057,11 @@ TEST(CApiTest, InsertSamePkAfterDeleteOnSealedSegment) { // insert data with pks = {0, 0, 1, 1, 2, 2, 3, 3, 4, 4} , timestamps = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9} auto segment_interface = reinterpret_cast(segment); auto sealed_segment = dynamic_cast(segment_interface); - SealedLoadFieldData(dataset, *sealed_segment); - sealed_segment->get_insert_record().seal_pks(); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareInsertBinlog( + kCollectionID, kPartitionID, kSegmentID, dataset, cm); + sealed_segment->LoadFieldData(load_info); // delete data pks = {1, 2, 3}, timestamps = {4, 4, 4} std::vector delete_row_ids = {1, 2, 3}; @@ -1082,10 +1071,7 @@ TEST(CApiTest, InsertSamePkAfterDeleteOnSealedSegment) { auto delete_data = serialize(ids.get()); std::vector delete_timestamps(3, dataset.timestamps_[4]); - auto offset = 0; - auto del_res = Delete(segment, - offset, 3, delete_data.data(), delete_data.size(), @@ -1372,14 +1358,8 @@ TEST(CApiTest, GetDeletedCountTest) { auto delete_data = serialize(ids.get()); uint64_t delete_timestamps[] = {0, 0, 0}; - auto offset = 0; - - auto del_res = Delete(segment, - offset, - 3, - delete_data.data(), - delete_data.size(), - delete_timestamps); + auto del_res = Delete( + segment, 3, delete_data.data(), delete_data.size(), delete_timestamps); ASSERT_EQ(del_res.error_code, Success); // TODO: assert(deleted_count == len(delete_row_ids)) @@ -1453,14 +1433,8 @@ TEST(CApiTest, GetRealCount) { dataset.timestamps_[N - 1] + 2, dataset.timestamps_[N - 1] + 3}; - auto del_offset = 0; - - auto del_res = Delete(segment, - del_offset, - 3, - delete_data.data(), - delete_data.size(), - delete_timestamps); + auto del_res = Delete( + segment, 3, delete_data.data(), delete_data.size(), delete_timestamps); ASSERT_EQ(del_res.error_code, Success); auto real_count = GetRealCount(segment); @@ -2097,7 +2071,7 @@ Test_Indexing_Without_Predicate() { knowhere::Version::GetCurrentVersion().VersionNumber()); AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -2254,7 +2228,7 @@ TEST(CApiTest, Indexing_Expr_Without_Predicate) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -2432,7 +2406,7 @@ TEST(CApiTest, Indexing_With_float_Predicate_Range) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -2612,7 +2586,7 @@ TEST(CApiTest, Indexing_Expr_With_float_Predicate_Range) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -2784,7 +2758,7 @@ TEST(CApiTest, Indexing_With_float_Predicate_Term) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -2957,7 +2931,7 @@ TEST(CApiTest, Indexing_Expr_With_float_Predicate_Term) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3138,7 +3112,7 @@ TEST(CApiTest, Indexing_With_binary_Predicate_Range) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3319,7 +3293,7 @@ TEST(CApiTest, Indexing_Expr_With_binary_Predicate_Range) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3494,7 +3468,7 @@ TEST(CApiTest, Indexing_With_binary_Predicate_Term) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3692,7 +3666,7 @@ TEST(CApiTest, Indexing_Expr_With_binary_Predicate_Term) { AppendIndex(c_load_index_info, (CBinarySet)&binary_set); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3741,27 +3715,6 @@ TEST(CApiTest, Indexing_Expr_With_binary_Predicate_Term) { DeleteSearchResultDataBlobs(cSearchResultData); } -TEST(CApiTest, SealedSegmentTest) { - auto collection = NewCollection(get_default_schema_config().c_str()); - CSegmentInterface segment; - auto status = NewSegment(collection, Sealed, -1, &segment, false); - ASSERT_EQ(status.error_code, Success); - - int N = 1000; - std::default_random_engine e(67); - auto ages = std::vector(N); - for (auto& age : ages) { - age = e() % 2000; - } - auto res = LoadFieldRawData(segment, 101, ages.data(), N); - ASSERT_EQ(res.error_code, Success); - auto count = GetRowCount(segment); - ASSERT_EQ(count, N); - - DeleteCollection(collection); - DeleteSegment(segment); -} - TEST(CApiTest, SealedSegment_search_float_Predicate_Range) { constexpr auto TOPK = 5; @@ -3778,8 +3731,6 @@ TEST(CApiTest, SealedSegment_search_float_Predicate_Range) { auto vec_col = dataset.get_col(FieldId(100)); auto query_ptr = vec_col.data() + BIAS * DIM; - auto counter_col = dataset.get_col(FieldId(101)); - const char* raw_plan = R"(vector_anns: < field_id: 100 predicates: < @@ -3880,17 +3831,21 @@ TEST(CApiTest, SealedSegment_search_float_Predicate_Range) { vec_index->Query(query_dataset, search_info, nullptr, result_on_index); EXPECT_EQ(result_on_index.distances_.size(), num_queries * TOPK); - status = LoadFieldRawData(segment, 101, counter_col.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 0, dataset.row_ids_.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 1, dataset.timestamps_.data(), N); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto excluded_field_ids = GetExcludedFieldIds(dataset.schema_, {0, 1, 101}); + auto load_info = PrepareInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + dataset, + cm, + "", + excluded_field_ids); + status = LoadFieldData(segment, &load_info); ASSERT_EQ(status.error_code, Success); // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(FieldId(100)); sealed_segment->LoadIndex(*(LoadIndexInfo*)c_load_index_info); @@ -3917,7 +3872,6 @@ TEST(CApiTest, SealedSegment_search_float_Predicate_Range) { } TEST(CApiTest, SealedSegment_search_without_predicates) { - constexpr auto TOPK = 5; std::string schema_string = generate_collection_schema( knowhere::metric::L2, DIM); auto collection = NewCollection(schema_string.c_str()); @@ -3926,16 +3880,8 @@ TEST(CApiTest, SealedSegment_search_without_predicates) { auto status = NewSegment(collection, Sealed, -1, &segment, false); ASSERT_EQ(status.error_code, Success); - auto N = ROW_COUNT; uint64_t ts_offset = 1000; - auto dataset = DataGen(schema, N, ts_offset); - auto vec_col = dataset.get_col(FieldId(100)); - auto query_ptr = vec_col.data() + BIAS * DIM; - - auto vec_array = dataset.get_col(FieldId(100)); - auto vec_data = serialize(vec_array.get()); - - auto counter_col = dataset.get_col(FieldId(101)); + auto dataset = DataGen(schema, ROW_COUNT, ts_offset); const char* raw_plan = R"(vector_anns: < field_id: 100 @@ -3949,16 +3895,18 @@ TEST(CApiTest, SealedSegment_search_without_predicates) { >)"; auto plan_str = translate_text_plan_to_binary_plan(raw_plan); - status = LoadFieldRawData(segment, 100, vec_data.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 101, counter_col.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 0, dataset.row_ids_.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 1, dataset.timestamps_.data(), N); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto excluded_field_ids = + GetExcludedFieldIds(dataset.schema_, {0, 1, 100, 101}); + auto load_info = PrepareInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + dataset, + cm, + "", + excluded_field_ids); + status = LoadFieldData(segment, &load_info); ASSERT_EQ(status.error_code, Success); int num_queries = 10; @@ -3977,14 +3925,17 @@ TEST(CApiTest, SealedSegment_search_without_predicates) { std::vector placeholderGroups; placeholderGroups.push_back(placeholderGroup); CSearchResult search_result; - auto res = - CSearch(segment, plan, placeholderGroup, N + ts_offset, &search_result); + auto res = CSearch( + segment, plan, placeholderGroup, ROW_COUNT + ts_offset, &search_result); std::cout << res.error_msg << std::endl; ASSERT_EQ(res.error_code, Success); CSearchResult search_result2; - auto res2 = CSearch( - segment, plan, placeholderGroup, N + ts_offset, &search_result2); + auto res2 = CSearch(segment, + plan, + placeholderGroup, + ROW_COUNT + ts_offset, + &search_result2); ASSERT_EQ(res2.error_code, Success); DeleteSearchPlan(plan); @@ -4002,9 +3953,6 @@ TEST(CApiTest, SealedSegment_search_float_With_Expr_Predicate_Range) { knowhere::metric::L2, DIM); auto collection = NewCollection(schema_string.c_str()); auto schema = ((segcore::Collection*)collection)->get_schema(); - CSegmentInterface segment; - auto status = NewSegment(collection, Sealed, -1, &segment, false); - ASSERT_EQ(status.error_code, Success); auto N = ROW_COUNT; auto dataset = DataGen(schema, N); @@ -4062,7 +4010,7 @@ TEST(CApiTest, SealedSegment_search_float_With_Expr_Predicate_Range) { // search on segment's small index void* plan = nullptr; auto binary_plan = translate_text_plan_to_binary_plan(serialized_expr_plan); - status = CreateSearchPlanByExpr( + auto status = CreateSearchPlanByExpr( collection, binary_plan.data(), binary_plan.size(), &plan); ASSERT_EQ(status.error_code, Success); @@ -4103,18 +4051,10 @@ TEST(CApiTest, SealedSegment_search_float_With_Expr_Predicate_Range) { knowhere::Version::GetCurrentVersion().VersionNumber()); AppendIndex(c_load_index_info, (CBinarySet)&binary_set); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); + // load vec index - status = UpdateSealedSegmentIndex(segment, c_load_index_info); - ASSERT_EQ(status.error_code, Success); - - // load raw data - status = LoadFieldRawData(segment, 101, counter_col.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 0, dataset.row_ids_.data(), N); - ASSERT_EQ(status.error_code, Success); - - status = LoadFieldRawData(segment, 1, dataset.timestamps_.data(), N); + status = UpdateSealedSegmentIndex(segment.get(), c_load_index_info); ASSERT_EQ(status.error_code, Success); // gen query dataset @@ -4133,7 +4073,7 @@ TEST(CApiTest, SealedSegment_search_float_With_Expr_Predicate_Range) { } CSearchResult c_search_result_on_bigIndex; - auto res_after_load_index = CSearch(segment, + auto res_after_load_index = CSearch(segment.get(), plan, placeholderGroup, timestamp, @@ -4151,39 +4091,6 @@ TEST(CApiTest, SealedSegment_search_float_With_Expr_Predicate_Range) { DeletePlaceholderGroup(placeholderGroup); DeleteSearchResult(c_search_result_on_bigIndex); DeleteCollection(collection); - DeleteSegment(segment); -} - -TEST(CApiTest, SealedSegment_Update_Field_Size) { - auto schema = std::make_shared(); - auto str_fid = schema->AddDebugField("string", DataType::VARCHAR); - auto vec_fid = schema->AddDebugField( - "vector_float", DataType::VECTOR_FLOAT, DIM, "L2"); - schema->set_primary_field_id(str_fid); - - auto segment = CreateSealedSegment(schema).release(); - - auto N = ROW_COUNT; - int row_size = 10; - - // update row_size =10 with n rows - auto status = - UpdateFieldRawDataSize(segment, str_fid.get(), N, N * row_size); - ASSERT_EQ(status.error_code, Success); - ASSERT_EQ(segment->get_field_avg_size(str_fid), row_size); - - // load data and update avg field size - std::vector str_datas; - int64_t total_size = 0; - for (int i = 0; i < N; ++i) { - auto str = "string_data_" + std::to_string(i); - total_size += str.size() + sizeof(uint32_t); - str_datas.emplace_back(str); - } - auto res = LoadFieldRawData(segment, str_fid.get(), str_datas.data(), N); - ASSERT_EQ(res.error_code, Success); - - DeleteSegment(segment); } TEST(CApiTest, GrowingSegment_Load_Field_Data) { @@ -4207,13 +4114,7 @@ TEST(CApiTest, GrowingSegment_Load_Field_Data) { auto storage_config = get_default_local_storage_config(); auto cm = storage::CreateChunkManager(storage_config); - auto load_info = - PrepareInsertBinlog(1, - 2, - 3, - storage_config.root_path + "/" + "test_load_sealed", - raw_data, - cm); + auto load_info = PrepareInsertBinlog(1, 2, 3, raw_data, cm); auto status = LoadFieldData(segment, &load_info); ASSERT_EQ(status.error_code, Success); @@ -4266,13 +4167,7 @@ TEST(CApiTest, GrowingSegment_Load_Field_Data_Lack_Binlog_Rows) { auto storage_config = get_default_local_storage_config(); auto cm = storage::CreateChunkManager(storage_config); - auto load_info = - PrepareInsertBinlog(1, - 2, - 3, - storage_config.root_path + "/" + "test_load_sealed", - raw_data, - cm); + auto load_info = PrepareInsertBinlog(1, 2, 3, raw_data, cm); raw_data.raw_->mutable_fields_data()->AddAllocated(array.release()); load_info.field_infos.emplace( @@ -4342,13 +4237,7 @@ TEST(CApiTest, DISABLED_SealedSegment_Load_Field_Data_Lack_Binlog_Rows) { auto storage_config = get_default_local_storage_config(); auto cm = storage::CreateChunkManager(storage_config); - auto load_info = - PrepareInsertBinlog(1, - 2, - 3, - storage_config.root_path + "/" + "test_load_sealed", - raw_data, - cm); + auto load_info = PrepareInsertBinlog(1, 2, 3, raw_data, cm); raw_data.raw_->mutable_fields_data()->AddAllocated(array.release()); load_info.field_infos.emplace( @@ -4375,7 +4264,7 @@ TEST(CApiTest, DISABLED_SealedSegment_Load_Field_Data_Lack_Binlog_Rows) { DeleteSegment(segment); } -TEST(CApiTest, RetriveScalarFieldFromSealedSegmentWithIndex) { +TEST(CApiTest, RetrieveScalarFieldFromSealedSegmentWithIndex) { auto schema = std::make_shared(); auto i8_fid = schema->AddDebugField("age8", DataType::INT8); auto i16_fid = schema->AddDebugField("age16", DataType::INT16); @@ -4391,26 +4280,19 @@ TEST(CApiTest, RetriveScalarFieldFromSealedSegmentWithIndex) { auto raw_data = DataGen(schema, N); LoadIndexInfo load_index_info; - // load timestamp field - auto res = LoadFieldRawData( - segment, TimestampFieldID.get(), raw_data.timestamps_.data(), N); - ASSERT_EQ(res.error_code, Success); - auto count = GetRowCount(segment); - ASSERT_EQ(count, N); - - // load rowid field - res = LoadFieldRawData( - segment, RowFieldID.get(), raw_data.row_ids_.data(), N); - ASSERT_EQ(res.error_code, Success); - count = GetRowCount(segment); - ASSERT_EQ(count, N); - - // load int64 field - res = LoadFieldRawData( - segment, i64_fid.get(), raw_data.get_col(i64_fid).data(), N); - ASSERT_EQ(res.error_code, Success); - count = GetRowCount(segment); - ASSERT_EQ(count, N); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto excluded_field_ids = + GetExcludedFieldIds(raw_data.schema_, {0, 1, i64_fid.get()}); + auto load_info = PrepareInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + raw_data, + cm, + "", + excluded_field_ids); + auto status = LoadFieldData(segment, &load_info); + ASSERT_EQ(status.error_code, Success); // load index for int8 field auto age8_col = raw_data.get_col(i8_fid); @@ -4492,7 +4374,7 @@ TEST(CApiTest, RetriveScalarFieldFromSealedSegmentWithIndex) { plan->field_ids_ = target_field_ids; CRetrieveResult* retrieve_result = nullptr; - res = CRetrieve( + auto res = CRetrieve( segment, plan.get(), raw_data.timestamps_[N - 1], &retrieve_result); ASSERT_EQ(res.error_code, Success); auto query_result = std::make_unique(); diff --git a/internal/core/unittest/test_cached_search_iterator.cpp b/internal/core/unittest/test_cached_search_iterator.cpp index 8867328a54..f99a7e7873 100644 --- a/internal/core/unittest/test_cached_search_iterator.cpp +++ b/internal/core/unittest/test_cached_search_iterator.cpp @@ -28,6 +28,7 @@ #include "segcore/InsertRecord.h" #include "mmap/ChunkedColumn.h" #include "test_utils/DataGen.h" +#include "test_cachinglayer/cachinglayer_test_utils.h" using namespace milvus; using namespace milvus::query; @@ -103,6 +104,8 @@ class CachedSearchIteratorTest static std::unique_ptr> vector_base_; static std::shared_ptr column_; static std::vector> column_data_; + static std::shared_ptr schema_; + static FieldId fakevec_id_; IndexBase* index_hnsw_ = nullptr; MetricType metric_type_ = kMetricType; @@ -140,7 +143,7 @@ class CachedSearchIteratorTest case ConstructorType::ChunkedColumn: return std::make_unique( - column_, + column_.get(), search_dataset_, search_info, std::map{}, @@ -245,35 +248,42 @@ class CachedSearchIteratorTest static void SetUpChunkedColumn() { - column_ = std::make_unique(); + auto field_meta = schema_->operator[](fakevec_id_); const size_t num_chunks_ = (nb_ + kSizePerChunk - 1) / kSizePerChunk; column_data_.resize(num_chunks_); size_t offset = 0; + std::vector> chunks; + std::vector num_rows_per_chunk; for (size_t i = 0; i < num_chunks_; ++i) { const size_t rows = std::min(static_cast(nb_ - offset), kSizePerChunk); + num_rows_per_chunk.push_back(rows); const size_t buf_size = rows * dim_ * sizeof(float); auto& chunk_data = column_data_[i]; chunk_data.resize(buf_size); memcpy(chunk_data.data(), base_dataset_.cbegin() + offset * dim_, rows * dim_ * sizeof(float)); - column_->AddChunk(std::make_shared( + chunks.emplace_back(std::make_unique( rows, dim_, chunk_data.data(), buf_size, sizeof(float), false)); offset += rows; } + auto translator = std::make_unique( + num_rows_per_chunk, "", std::move(chunks)); + column_ = + std::make_shared(std::move(translator), field_meta); } static void SetUpTestSuite() { - auto schema = std::make_shared(); - auto fakevec_id = schema->AddDebugField( + schema_ = std::make_shared(); + fakevec_id_ = schema_->AddDebugField( "fakevec", DataType::VECTOR_FLOAT, dim_, kMetricType); // generate base dataset base_dataset_ = - segcore::DataGen(schema, nb_).get_col(fakevec_id); + segcore::DataGen(schema_, nb_).get_col(fakevec_id_); // generate query dataset query_dataset_ = {base_dataset_.cbegin(), @@ -348,6 +358,8 @@ std::unique_ptr> CachedSearchIteratorTest::vector_base_ = nullptr; std::shared_ptr CachedSearchIteratorTest::column_ = nullptr; std::vector> CachedSearchIteratorTest::column_data_; +std::shared_ptr CachedSearchIteratorTest::schema_{nullptr}; +FieldId CachedSearchIteratorTest::fakevec_id_(0); /********* Testcases Start **********/ @@ -605,7 +617,6 @@ TEST_P(CachedSearchIteratorTest, NextBatchtAllBatchesNormal) { SearchInfo search_info = GetDefaultNormalSearchInfo(); const std::vector kBatchSizes = { 1, 7, 43, 99, 100, 101, 1000, 1005}; - // const std::vector kBatchSizes = {1005}; for (size_t batch_size : kBatchSizes) { search_info.iterator_v2_info_->batch_size = batch_size; @@ -745,7 +756,7 @@ TEST_P(CachedSearchIteratorTest, ConstructorWithInvalidParams) { data_type_), SegcoreError); EXPECT_THROW(auto iterator = std::make_unique( - column_, + column_.get(), dataset::SearchDataset{}, search_info, std::map{}, diff --git a/internal/core/unittest/test_cachinglayer/CMakeLists.txt b/internal/core/unittest/test_cachinglayer/CMakeLists.txt new file mode 100644 index 0000000000..78a6dcf4b9 --- /dev/null +++ b/internal/core/unittest/test_cachinglayer/CMakeLists.txt @@ -0,0 +1,29 @@ +# Copyright (C) 2019-2020 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under the License + +set(CACHINGLAYER_TEST_FILES + ../init_gtest.cpp + test_dlist.cpp + test_cache_slot.cpp +) + +add_executable(cachinglayer_test + ${CACHINGLAYER_TEST_FILES} +) + +target_link_libraries(cachinglayer_test + gtest + gmock + milvus_core + milvus-storage +) + +install(TARGETS cachinglayer_test DESTINATION unittest) diff --git a/internal/core/unittest/test_cachinglayer/cachinglayer_test_utils.h b/internal/core/unittest/test_cachinglayer/cachinglayer_test_utils.h new file mode 100644 index 0000000000..205892c35f --- /dev/null +++ b/internal/core/unittest/test_cachinglayer/cachinglayer_test_utils.h @@ -0,0 +1,171 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include + +#include +#include + +#include "cachinglayer/Translator.h" +#include "common/Chunk.h" +#include "segcore/storagev1translator/ChunkTranslator.h" +#include "cachinglayer/lrucache/DList.h" + +namespace milvus { + +using namespace cachinglayer; + +class TestChunkTranslator : public Translator { + public: + TestChunkTranslator(std::vector num_rows_per_chunk, + std::string key, + std::vector>&& chunks) + : Translator(), + num_cells_(num_rows_per_chunk.size()), + chunks_(std::move(chunks)), + meta_(segcore::storagev1translator::CTMeta(StorageType::MEMORY)) { + meta_.num_rows_until_chunk_.reserve(num_cells_ + 1); + meta_.num_rows_until_chunk_.push_back(0); + int total_rows = 0; + for (int i = 0; i < num_cells_; ++i) { + meta_.num_rows_until_chunk_.push_back( + meta_.num_rows_until_chunk_[i] + num_rows_per_chunk[i]); + total_rows += num_rows_per_chunk[i]; + } + key_ = key; + } + ~TestChunkTranslator() override { + } + + size_t + num_cells() const override { + return num_cells_; + } + + cid_t + cell_id_of(uid_t uid) const override { + return uid; + } + + ResourceUsage + estimated_byte_size_of_cell(cid_t cid) const override { + return ResourceUsage(0, 0); + } + + const std::string& + key() const override { + return key_; + } + + Meta* + meta() override { + return &meta_; + } + + std::vector>> + get_cells(const std::vector& cids) override { + std::vector>> res; + res.reserve(cids.size()); + for (auto cid : cids) { + AssertInfo(cid < chunks_.size() && chunks_[cid] != nullptr, + "TestChunkTranslator assumes no eviction."); + res.emplace_back(cid, std::move(chunks_[cid])); + } + return res; + } + + private: + size_t num_cells_; + segcore::storagev1translator::CTMeta meta_; + std::string key_; + std::vector> chunks_; +}; + +namespace cachinglayer::internal { +class DListTestFriend { + public: + static ResourceUsage + get_used_memory(const DList& dlist) { + return dlist.used_memory_.load(); + } + static ResourceUsage + get_max_memory(const DList& dlist) { + std::lock_guard lock(dlist.list_mtx_); + return dlist.max_memory_; + } + static ListNode* + get_head(const DList& dlist) { + std::lock_guard lock(dlist.list_mtx_); + return dlist.head_; + } + static ListNode* + get_tail(const DList& dlist) { + std::lock_guard lock(dlist.list_mtx_); + return dlist.tail_; + } + static void + test_push_head(DList* dlist, ListNode* node) { + std::lock_guard lock(dlist->list_mtx_); + dlist->pushHead(node); + } + static void + test_pop_item(DList* dlist, ListNode* node) { + std::lock_guard lock(dlist->list_mtx_); + dlist->popItem(node); + } + static void + test_add_used_memory(DList* dlist, const ResourceUsage& size) { + std::lock_guard lock(dlist->list_mtx_); + dlist->used_memory_ += size; + } + + // nodes are from tail to head + static void + verify_list(DList* dlist, std::vector nodes) { + std::lock_guard lock(dlist->list_mtx_); + EXPECT_EQ(nodes.front(), dlist->tail_); + EXPECT_EQ(nodes.back(), dlist->head_); + for (size_t i = 0; i < nodes.size(); ++i) { + auto current = nodes[i]; + auto expected_prev = i == 0 ? nullptr : nodes[i - 1]; + auto expected_next = i == nodes.size() - 1 ? nullptr : nodes[i + 1]; + EXPECT_EQ(current->prev_, expected_prev); + EXPECT_EQ(current->next_, expected_next); + } + } + + static void + verify_integrity(DList* dlist) { + std::lock_guard lock(dlist->list_mtx_); + + ResourceUsage total_size; + EXPECT_EQ(dlist->tail_->prev_, nullptr); + ListNode* current = dlist->tail_; + ListNode* prev = nullptr; + + while (current != nullptr) { + EXPECT_EQ(current->prev_, prev); + total_size += current->size(); + prev = current; + current = current->next_; + } + + EXPECT_EQ(prev, dlist->head_); + EXPECT_EQ(dlist->head_->next_, nullptr); + + EXPECT_EQ(total_size, dlist->used_memory_.load()); + } +}; +} // namespace cachinglayer::internal + +} // namespace milvus diff --git a/internal/core/unittest/test_cachinglayer/mock_list_node.h b/internal/core/unittest/test_cachinglayer/mock_list_node.h new file mode 100644 index 0000000000..2eecf00771 --- /dev/null +++ b/internal/core/unittest/test_cachinglayer/mock_list_node.h @@ -0,0 +1,85 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "cachinglayer/lrucache/ListNode.h" +#include "cachinglayer/Utils.h" + +namespace milvus::cachinglayer::internal { + +class MockListNode : public ListNode { + public: + MockListNode(DList* dlist, + ResourceUsage size, + const std::string& key = "mock_key", + cid_t cid = 0) + : ListNode(dlist, size), mock_key_(fmt::format("{}:{}", key, cid)) { + ON_CALL(*this, clear_data).WillByDefault([this]() { + // Default clear_data calls unload() by default in base, mimic if needed + unload(); + state_ = State::NOT_LOADED; + }); + } + + MOCK_METHOD(void, clear_data, (), (override)); + + std::string + key() const override { + return mock_key_; + } + + // Directly manipulate state for test setup (Use carefully!) + void + test_set_state(State new_state) { + std::unique_lock lock(mtx_); + state_ = new_state; + } + State + test_get_state() { + std::shared_lock lock(mtx_); + return state_; + } + + void + test_set_pin_count(int count) { + pin_count_.store(count); + } + int + test_get_pin_count() const { + return pin_count_.load(); + } + + // Expose mutex for lock testing + std::shared_mutex& + test_get_mutex() { + return mtx_; + } + + ListNode* + test_get_prev() const { + return prev_; + } + ListNode* + test_get_next() const { + return next_; + } + + private: + friend class DListTest; + friend class DListTestFriend; + std::string mock_key_; +}; + +} // namespace milvus::cachinglayer::internal \ No newline at end of file diff --git a/internal/core/unittest/test_cachinglayer/test_cache_slot.cpp b/internal/core/unittest/test_cachinglayer/test_cache_slot.cpp new file mode 100644 index 0000000000..d726fdcdfa --- /dev/null +++ b/internal/core/unittest/test_cachinglayer/test_cache_slot.cpp @@ -0,0 +1,846 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cachinglayer/CacheSlot.h" +#include "cachinglayer/Translator.h" +#include "cachinglayer/Utils.h" +#include "cachinglayer/lrucache/DList.h" +#include "cachinglayer/lrucache/ListNode.h" +#include "cachinglayer_test_utils.h" + +using namespace milvus::cachinglayer; +using namespace milvus::cachinglayer::internal; +using cl_uid_t = milvus::cachinglayer::uid_t; + +struct TestCell { + int data; + cid_t cid; + + TestCell(int d, cid_t id) : data(d), cid(id) { + } + + size_t + CellByteSize() const { + return sizeof(data) + sizeof(cid); + } +}; + +class MockTranslator : public Translator { + public: + MockTranslator(std::vector> cell_sizes, + std::unordered_map uid_to_cid_map, + const std::string& key, + StorageType storage_type, + bool for_concurrent_test = false) + : uid_to_cid_map_(std::move(uid_to_cid_map)), + key_(key), + meta_(storage_type), + num_unique_cids_(cell_sizes.size()) { + cid_set_.reserve(cell_sizes.size()); + cell_sizes_.reserve(cell_sizes.size()); + for (const auto& pair : cell_sizes) { + cid_t cid = pair.first; + int64_t size = pair.second; + cid_set_.insert(cid); + cell_sizes_[cid] = size; + cid_load_delay_ms_[cid] = 0; + } + } + + size_t + num_cells() const override { + return num_unique_cids_; + } + + cid_t + cell_id_of(cl_uid_t uid) const override { + auto it = uid_to_cid_map_.find(uid); + if (it != uid_to_cid_map_.end()) { + if (cid_set_.count(it->second)) { + return it->second; + } + } + return static_cast(num_unique_cids_); + } + + ResourceUsage + estimated_byte_size_of_cell(cid_t cid) const override { + auto it = cell_sizes_.find(cid); + if (it != cell_sizes_.end()) { + return ResourceUsage{it->second, 0}; + } + return ResourceUsage{1, 0}; + } + + const std::string& + key() const override { + return key_; + } + + Meta* + meta() override { + return &meta_; + } + + std::vector>> + get_cells(const std::vector& cids) override { + if (!for_concurrent_test_) { + get_cells_call_count_++; + requested_cids_.push_back(cids); + } + + if (load_should_throw_) { + throw std::runtime_error("Simulated load error"); + } + + std::vector>> result; + for (cid_t cid : cids) { + auto delay_it = cid_load_delay_ms_.find(cid); + if (delay_it != cid_load_delay_ms_.end() && delay_it->second > 0) { + std::this_thread::sleep_for( + std::chrono::milliseconds(delay_it->second)); + } + + result.emplace_back( + cid, + std::make_unique(static_cast(cid * 10), cid)); + + if (auto extra_cids = extra_cids_.find(cid); + extra_cids != extra_cids_.end()) { + for (cid_t extra_cid : extra_cids->second) { + // if extra cid is not explicitly requested, and not yet added as extra cell by other + // cells, add it to the result. + if (std::find(cids.begin(), cids.end(), extra_cid) == + cids.end() && + std::find_if(result.begin(), + result.end(), + [extra_cid](const auto& pair) { + return pair.first == extra_cid; + }) == result.end()) { + result.emplace_back( + extra_cid, + std::make_unique( + static_cast(extra_cid * 10), extra_cid)); + } + } + } + } + return result; + } + + void + SetCidLoadDelay(const std::unordered_map& delays) { + for (const auto& pair : delays) { + cid_load_delay_ms_[pair.first] = pair.second; + } + } + void + SetShouldThrow(bool should_throw) { + load_should_throw_ = should_throw; + } + // for some cid, translator will return extra cells. + void + SetExtraReturnCids( + std::unordered_map> extra_cids) { + extra_cids_ = extra_cids; + } + int + GetCellsCallCount() const { + EXPECT_FALSE(for_concurrent_test_); + return get_cells_call_count_; + } + const std::vector>& + GetRequestedCids() const { + EXPECT_FALSE(for_concurrent_test_); + return requested_cids_; + } + void + ResetCounters() { + ASSERT_FALSE(for_concurrent_test_); + get_cells_call_count_ = 0; + requested_cids_.clear(); + } + + private: + std::unordered_map uid_to_cid_map_; + std::unordered_map cell_sizes_; + std::unordered_set cid_set_; + const size_t num_unique_cids_; + const std::string key_; + Meta meta_; + + std::unordered_map cid_load_delay_ms_; + bool load_should_throw_ = false; + std::unordered_map> extra_cids_; + std::atomic get_cells_call_count_ = 0; + std::vector> requested_cids_; + + // this class is not concurrent safe, so if for concurrent test, do not track usage + bool for_concurrent_test_ = false; +}; + +class CacheSlotTest : public ::testing::Test { + protected: + std::unique_ptr dlist_; + MockTranslator* translator_ = nullptr; + std::shared_ptr> cache_slot_; + + std::vector> cell_sizes_ = { + {0, 50}, {1, 150}, {2, 100}, {3, 200}, {4, 75}}; + std::unordered_map uid_to_cid_map_ = {{10, 0}, + {11, 0}, + {20, 1}, + {30, 2}, + {31, 2}, + {32, 2}, + {40, 3}, + {50, 4}, + {51, 4}}; + + size_t NUM_UNIQUE_CIDS = 5; + int64_t TOTAL_CELL_SIZE_BYTES = 50 + 150 + 100 + 200 + 75; + int64_t MEMORY_LIMIT = TOTAL_CELL_SIZE_BYTES * 2; + static constexpr int64_t DISK_LIMIT = 0; + const std::string SLOT_KEY = "test_slot"; + + void + SetUp() override { + dlist_ = std::make_unique( + ResourceUsage{MEMORY_LIMIT, DISK_LIMIT}, DList::TouchConfig{}); + + auto temp_translator_uptr = std::make_unique( + cell_sizes_, uid_to_cid_map_, SLOT_KEY, StorageType::MEMORY); + translator_ = temp_translator_uptr.get(); + cache_slot_ = std::make_shared>( + std::move(temp_translator_uptr), dlist_.get()); + } + + void + TearDown() override { + cache_slot_.reset(); + dlist_.reset(); + } +}; + +TEST_F(CacheSlotTest, Initialization) { + ASSERT_EQ(cache_slot_->num_cells(), NUM_UNIQUE_CIDS); +} + +TEST_F(CacheSlotTest, PinSingleCellSuccess) { + cl_uid_t target_uid = 30; + cid_t expected_cid = 2; + ResourceUsage expected_size = + translator_->estimated_byte_size_of_cell(expected_cid); + + translator_->ResetCounters(); + auto future = cache_slot_->PinCells({target_uid}); + auto accessor = SemiInlineGet(std::move(future)); + + ASSERT_NE(accessor, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + ASSERT_EQ(translator_->GetRequestedCids()[0].size(), 1); + EXPECT_EQ(translator_->GetRequestedCids()[0][0], expected_cid); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + + TestCell* cell = accessor->get_cell_of(target_uid); + ASSERT_NE(cell, nullptr); + EXPECT_EQ(cell->cid, expected_cid); + EXPECT_EQ(cell->data, expected_cid * 10); + + TestCell* cell_by_index = accessor->get_ith_cell(expected_cid); + ASSERT_EQ(cell, cell_by_index); +} + +TEST_F(CacheSlotTest, PinMultipleCellsSuccess) { + std::vector target_uids = {10, 40, 51}; + std::vector expected_cids = {0, 3, 4}; + std::sort(expected_cids.begin(), expected_cids.end()); + ResourceUsage expected_total_size; + for (cid_t cid : expected_cids) { + expected_total_size += translator_->estimated_byte_size_of_cell(cid); + } + + translator_->ResetCounters(); + auto future = cache_slot_->PinCells(target_uids); + auto accessor = SemiInlineGet(std::move(future)); + + ASSERT_NE(accessor, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + auto requested = translator_->GetRequestedCids()[0]; + std::sort(requested.begin(), requested.end()); + ASSERT_EQ(requested.size(), expected_cids.size()); + EXPECT_EQ(requested, expected_cids); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_total_size); + + for (cl_uid_t uid : target_uids) { + cid_t cid = uid_to_cid_map_.at(uid); + TestCell* cell = accessor->get_cell_of(uid); + ASSERT_NE(cell, nullptr); + EXPECT_EQ(cell->cid, cid); + EXPECT_EQ(cell->data, cid * 10); + } +} + +TEST_F(CacheSlotTest, PinMultipleUidsMappingToSameCid) { + std::vector target_uids = {30, 50, 31, 51, 32}; + std::vector expected_unique_cids = {2, 4}; + std::sort(expected_unique_cids.begin(), expected_unique_cids.end()); + ResourceUsage expected_total_size; + for (cid_t cid : expected_unique_cids) { + expected_total_size += translator_->estimated_byte_size_of_cell(cid); + } + + translator_->ResetCounters(); + auto future = cache_slot_->PinCells(target_uids); + auto accessor = SemiInlineGet(std::move(future)); + + ASSERT_NE(accessor, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + auto requested = translator_->GetRequestedCids()[0]; + std::sort(requested.begin(), requested.end()); + ASSERT_EQ(requested.size(), expected_unique_cids.size()); + EXPECT_EQ(requested, expected_unique_cids); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_total_size); + + TestCell* cell2_uid30 = accessor->get_cell_of(30); + TestCell* cell2_uid31 = accessor->get_cell_of(31); + TestCell* cell4_uid50 = accessor->get_cell_of(50); + TestCell* cell4_uid51 = accessor->get_cell_of(51); + ASSERT_NE(cell2_uid30, nullptr); + ASSERT_NE(cell4_uid50, nullptr); + EXPECT_EQ(cell2_uid30->cid, 2); + EXPECT_EQ(cell4_uid50->cid, 4); + EXPECT_EQ(cell2_uid30, cell2_uid31); + EXPECT_EQ(cell4_uid50, cell4_uid51); +} + +TEST_F(CacheSlotTest, PinInvalidUid) { + cl_uid_t invalid_uid = 999; + cl_uid_t valid_uid = 10; + std::vector target_uids = {valid_uid, invalid_uid}; + + translator_->ResetCounters(); + auto future = cache_slot_->PinCells(target_uids); + + EXPECT_THROW( + { + try { + SemiInlineGet(std::move(future)); + } catch (const std::invalid_argument& e) { + std::string error_what = e.what(); + EXPECT_TRUE(error_what.find("out of range") != + std::string::npos || + error_what.find("invalid") != std::string::npos); + throw; + } + }, + std::invalid_argument); + + EXPECT_EQ(translator_->GetCellsCallCount(), 0); +} + +TEST_F(CacheSlotTest, LoadFailure) { + cl_uid_t target_uid = 20; + cid_t expected_cid = 1; + + translator_->ResetCounters(); + translator_->SetShouldThrow(true); + + auto future = cache_slot_->PinCells({target_uid}); + + EXPECT_THROW( + { + try { + SemiInlineGet(std::move(future)); + } catch (const std::runtime_error& e) { + std::string error_what = e.what(); + EXPECT_TRUE(error_what.find("Simulated load error") != + std::string::npos || + error_what.find("Failed to load") != + std::string::npos || + error_what.find("Exception during Future") != + std::string::npos); + throw; + } + }, + std::runtime_error); + + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + ASSERT_EQ(translator_->GetRequestedCids()[0].size(), 1); + EXPECT_EQ(translator_->GetRequestedCids()[0][0], expected_cid); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), ResourceUsage{}); +} + +TEST_F(CacheSlotTest, PinAlreadyLoadedCell) { + cl_uid_t target_uid = 40; + cid_t expected_cid = 3; + ResourceUsage expected_size = + translator_->estimated_byte_size_of_cell(expected_cid); + + translator_->ResetCounters(); + + auto future1 = cache_slot_->PinCells({target_uid}); + auto accessor1 = SemiInlineGet(std::move(future1)); + ASSERT_NE(accessor1, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + ASSERT_EQ(translator_->GetRequestedCids()[0][0], expected_cid); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + TestCell* cell1 = accessor1->get_cell_of(target_uid); + ASSERT_NE(cell1, nullptr); + + translator_->ResetCounters(); + auto future2 = cache_slot_->PinCells({target_uid}); + auto accessor2 = SemiInlineGet(std::move(future2)); + ASSERT_NE(accessor2, nullptr); + + EXPECT_EQ(translator_->GetCellsCallCount(), 0); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + + TestCell* cell2 = accessor2->get_cell_of(target_uid); + ASSERT_NE(cell2, nullptr); + EXPECT_EQ(cell1, cell2); + + accessor1.reset(); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + TestCell* cell_after_unpin = accessor2->get_cell_of(target_uid); + ASSERT_NE(cell_after_unpin, nullptr); + EXPECT_EQ(cell_after_unpin, cell2); +} + +TEST_F(CacheSlotTest, PinAlreadyLoadedCellViaDifferentUid) { + cl_uid_t uid1 = 30; + cl_uid_t uid2 = 31; + cid_t expected_cid = 2; + ResourceUsage expected_size = + translator_->estimated_byte_size_of_cell(expected_cid); + + translator_->ResetCounters(); + + auto future1 = cache_slot_->PinCells({uid1}); + auto accessor1 = SemiInlineGet(std::move(future1)); + ASSERT_NE(accessor1, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + ASSERT_EQ(translator_->GetRequestedCids()[0][0], expected_cid); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + TestCell* cell1 = accessor1->get_cell_of(uid1); + ASSERT_NE(cell1, nullptr); + EXPECT_EQ(cell1->cid, expected_cid); + + translator_->ResetCounters(); + auto future2 = cache_slot_->PinCells({uid2}); + auto accessor2 = SemiInlineGet(std::move(future2)); + ASSERT_NE(accessor2, nullptr); + + EXPECT_EQ(translator_->GetCellsCallCount(), 0); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + + TestCell* cell2 = accessor2->get_cell_of(uid2); + ASSERT_NE(cell2, nullptr); + EXPECT_EQ(cell2->cid, expected_cid); + EXPECT_EQ(cell1, cell2); + + accessor1.reset(); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + TestCell* cell_after_unpin_uid1 = accessor2->get_cell_of(uid1); + TestCell* cell_after_unpin_uid2 = accessor2->get_cell_of(uid2); + ASSERT_NE(cell_after_unpin_uid1, nullptr); + ASSERT_NE(cell_after_unpin_uid2, nullptr); + EXPECT_EQ(cell_after_unpin_uid1, cell2); + EXPECT_EQ(cell_after_unpin_uid2, cell2); +} + +TEST_F(CacheSlotTest, TranslatorReturnsExtraCells) { + cl_uid_t requested_uid = 10; + cid_t requested_cid = 0; + cid_t extra_cid = 1; + cl_uid_t extra_uid = 20; + + ResourceUsage expected_size = + translator_->estimated_byte_size_of_cell(requested_cid) + + translator_->estimated_byte_size_of_cell(extra_cid); + + translator_->ResetCounters(); + translator_->SetExtraReturnCids({{requested_cid, {extra_cid}}}); + + auto future = cache_slot_->PinCells({requested_uid}); + auto accessor = SemiInlineGet(std::move(future)); + + ASSERT_NE(accessor, nullptr); + ASSERT_EQ(translator_->GetCellsCallCount(), 1); + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + EXPECT_EQ(translator_->GetRequestedCids()[0], + std::vector{requested_cid}); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + + TestCell* requested_cell = accessor->get_cell_of(requested_uid); + ASSERT_NE(requested_cell, nullptr); + EXPECT_EQ(requested_cell->cid, requested_cid); + + translator_->ResetCounters(); + auto future_extra = cache_slot_->PinCells({extra_uid}); + auto accessor_extra = SemiInlineGet(std::move(future_extra)); + + ASSERT_NE(accessor_extra, nullptr); + EXPECT_EQ(translator_->GetCellsCallCount(), 0); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), expected_size); + + TestCell* extra_cell = accessor_extra->get_cell_of(extra_uid); + ASSERT_NE(extra_cell, nullptr); + EXPECT_EQ(extra_cell->cid, extra_cid); +} + +TEST_F(CacheSlotTest, EvictionTest) { + // Sizes: 0:50, 1:150, 2:100, 3:200 + ResourceUsage NEW_LIMIT = ResourceUsage(300, 0); + EXPECT_TRUE(dlist_->UpdateLimit(NEW_LIMIT)); + EXPECT_EQ(DListTestFriend::get_max_memory(*dlist_), NEW_LIMIT); + + std::vector uids_012 = {10, 20, 30}; + std::vector cids_012 = {0, 1, 2}; + ResourceUsage size_012 = translator_->estimated_byte_size_of_cell(0) + + translator_->estimated_byte_size_of_cell(1) + + translator_->estimated_byte_size_of_cell(2); + ASSERT_EQ(size_012, ResourceUsage(50 + 150 + 100, 0)); + + // 1. Load cells 0, 1, 2 + translator_->ResetCounters(); + auto future1 = cache_slot_->PinCells(uids_012); + auto accessor1 = SemiInlineGet(std::move(future1)); + ASSERT_NE(accessor1, nullptr); + EXPECT_EQ(translator_->GetCellsCallCount(), 1); + auto requested1 = translator_->GetRequestedCids()[0]; + std::sort(requested1.begin(), requested1.end()); + EXPECT_EQ(requested1, cids_012); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), size_012); + + // 2. Unpin 0, 1, 2 + accessor1.reset(); + EXPECT_EQ(DListTestFriend::get_used_memory(*dlist_), + size_012); // Still in cache + + // 3. Load cell 3 (size 200), requires eviction + cl_uid_t uid_3 = 40; + cid_t cid_3 = 3; + ResourceUsage size_3 = translator_->estimated_byte_size_of_cell(cid_3); + ASSERT_EQ(size_3, ResourceUsage(200, 0)); + + translator_->ResetCounters(); + auto future2 = cache_slot_->PinCells({uid_3}); + auto accessor2 = SemiInlineGet(std::move(future2)); + ASSERT_NE(accessor2, nullptr); + + EXPECT_EQ(translator_->GetCellsCallCount(), + 1); // Load was called for cell 3 + ASSERT_EQ(translator_->GetRequestedCids().size(), 1); + EXPECT_EQ(translator_->GetRequestedCids()[0], std::vector{cid_3}); + + // Verify eviction happened + ResourceUsage used_after_evict1 = DListTestFriend::get_used_memory(*dlist_); + EXPECT_LE(used_after_evict1.memory_bytes, NEW_LIMIT.memory_bytes); + EXPECT_GE(used_after_evict1.memory_bytes, size_3.memory_bytes); + EXPECT_LT( + used_after_evict1.memory_bytes, + size_012.memory_bytes + size_3.memory_bytes); // Eviction occurred +} + +class CacheSlotConcurrentTest : public CacheSlotTest, + public ::testing::WithParamInterface {}; + +TEST_P(CacheSlotConcurrentTest, ConcurrentAccessMultipleSlots) { + // Slot 1 Cells: 0-4 (Sizes: 50, 60, 70, 80, 90) -> Total 350 + // Slot 2 Cells: 0-4 (Sizes: 55, 65, 75, 85, 95) -> Total 375 + // Total potential size = 350 + 375 = 725 + // Set limit lower than total potential size to force eviction + ResourceUsage NEW_LIMIT = ResourceUsage(600, 0); + ASSERT_TRUE(dlist_->UpdateLimit(NEW_LIMIT)); + EXPECT_EQ(DListTestFriend::get_max_memory(*dlist_).memory_bytes, + NEW_LIMIT.memory_bytes); + + // 1. Setup CacheSlots sharing dlist_ + std::vector> cell_sizes_1 = { + {0, 50}, {1, 60}, {2, 70}, {3, 80}, {4, 90}}; + std::unordered_map uid_map_1 = { + {1000, 0}, {1001, 1}, {1002, 2}, {1003, 3}, {1004, 4}}; + auto translator_1_ptr = std::make_unique( + cell_sizes_1, uid_map_1, "slot1", StorageType::MEMORY, true); + MockTranslator* translator_1 = translator_1_ptr.get(); + auto slot1 = std::make_shared>( + std::move(translator_1_ptr), dlist_.get()); + + std::vector> cell_sizes_2 = { + {0, 55}, {1, 65}, {2, 75}, {3, 85}, {4, 95}}; + std::unordered_map uid_map_2 = { + {2000, 0}, {2001, 1}, {2002, 2}, {2003, 3}, {2004, 4}}; + auto translator_2_ptr = std::make_unique( + cell_sizes_2, uid_map_2, "slot2", StorageType::MEMORY, true); + MockTranslator* translator_2 = translator_2_ptr.get(); + auto slot2 = std::make_shared>( + std::move(translator_2_ptr), dlist_.get()); + + bool with_bonus_cells = GetParam(); + if (with_bonus_cells) { + // Configure translators to return bonus cells + std::unordered_map> bonus_map_1{ + {0, {2}}, + {1, {3}}, + {2, {1, 4}}, + {3, {0}}, + {4, {2, 3}}, + }; + std::unordered_map> bonus_map_2{ + {0, {1, 4}}, + {1, {2, 3}}, + {2, {0}}, + {3, {0, 1}}, + {4, {2}}, + }; + translator_1->SetExtraReturnCids(bonus_map_1); + translator_2->SetExtraReturnCids(bonus_map_2); + } + + std::vector>> slots = {slot1, slot2}; + // Store uid maps in a structure easily accessible by slot index + std::vector> slot_uids; + slot_uids.resize(slots.size()); + std::vector> uid_to_cid_maps = { + uid_map_1, uid_map_2}; + for (const auto& pair : uid_map_1) slot_uids[0].push_back(pair.first); + for (const auto& pair : uid_map_2) slot_uids[1].push_back(pair.first); + + // 2. Setup Thread Pool and Concurrency Parameters + // at most 6 cells can be pinned at the same time, thus will never exceed the limit. + int num_threads = 6; + int ops_per_thread = 200; + // 1 extra thread to work with slot3 + folly::CPUThreadPoolExecutor executor(num_threads + 1); + std::vector> futures; + std::atomic test_failed{false}; + + // 3. Launch Threads to Perform Concurrent Pin/Get/Verify/Unpin + for (int i = 0; i < num_threads; ++i) { + futures.push_back(folly::via(&executor, [&, i, tid = i]() { + // Seed random generator uniquely for each thread + std::mt19937 gen( + std::hash{}(std::this_thread::get_id()) + tid); + std::uniform_int_distribution<> slot_dist(0, slots.size() - 1); + std::uniform_int_distribution<> sleep_dist(5, 15); + + for (int j = 0; j < ops_per_thread && !test_failed.load(); ++j) { + int slot_idx = slot_dist(gen); + auto& current_slot = slots[slot_idx]; + auto& current_slot_uids = slot_uids[slot_idx]; + auto& current_uid_to_cid_map = uid_to_cid_maps[slot_idx]; + + std::uniform_int_distribution<> uid_idx_dist( + 0, current_slot_uids.size() - 1); + cl_uid_t target_uid = current_slot_uids[uid_idx_dist(gen)]; + cid_t expected_cid = current_uid_to_cid_map.at(target_uid); + int expected_data = static_cast(expected_cid * 10); + + try { + auto accessor = current_slot->PinCells({target_uid}).get(); + + if (!accessor) { + ADD_FAILURE() + << "T" << tid << " Op" << j + << ": PinCells returned null accessor for UID " + << target_uid; + test_failed = true; + break; + } + + TestCell* cell = accessor->get_cell_of(target_uid); + if (!cell) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": get_cell_of returned null for UID " + << target_uid; + test_failed = true; + break; + } + + if (cell->cid != expected_cid) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Incorrect CID for UID " + << target_uid << ". Slot: " << slot_idx + << ", Expected: " << expected_cid + << ", Got: " << cell->cid; + test_failed = true; + break; + } + if (cell->data != expected_data) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Incorrect Data for UID " + << target_uid << ". Slot: " << slot_idx + << ", Expected: " << expected_data + << ", Got: " << cell->data; + test_failed = true; + break; + } + int sleep_ms = sleep_dist(gen); + std::this_thread::sleep_for( + std::chrono::milliseconds(sleep_ms)); + } catch (const std::exception& e) { + ADD_FAILURE() + << "T" << tid << " Op" << j << ": Exception for UID " + << target_uid << ", Slot: " << slot_idx + << ". What: " << e.what(); + test_failed = true; + } catch (...) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Unknown exception for UID " + << target_uid << ", Slot: " << slot_idx; + test_failed = true; + } + } + })); + } + + // number of ops between recreating slot3 + const int recreate_interval = 25; + auto dlist_ptr = dlist_.get(); + std::vector> cell_sizes_3 = { + {0, 40}, {1, 50}, {2, 60}, {3, 70}, {4, 80}}; + std::unordered_map uid_map_3 = { + {3000, 0}, {3001, 1}, {3002, 2}, {3003, 3}, {3004, 4}}; + std::vector slot3_uids = {3000, 3001, 3002, 3003, 3004}; + auto create_new_slot3 = [&]() { + auto translator_3_ptr = std::make_unique( + cell_sizes_3, uid_map_3, "slot3", StorageType::MEMORY, true); + auto sl = std::make_shared>( + std::move(translator_3_ptr), dlist_ptr); + // std::cout << "Created new SSSSSSslot3 at " << (void*)sl.get() << std::endl; + return sl; + }; + std::shared_ptr> slot3 = create_new_slot3(); + futures.push_back(folly::via(&executor, [&, tid = num_threads]() { + std::mt19937 gen( + std::hash{}(std::this_thread::get_id()) + tid); + std::uniform_int_distribution<> sleep_dist(5, 15); + std::uniform_int_distribution<> recreate_sleep_dist(20, 30); + std::uniform_int_distribution<> uid_idx_dist(0, slot3_uids.size() - 1); + int ops_since_recreate = 0; + + for (int j = 0; j < ops_per_thread && !test_failed.load(); ++j) { + cl_uid_t target_uid = slot3_uids[uid_idx_dist(gen)]; + cid_t expected_cid = uid_map_3.at(target_uid); + int expected_data = static_cast(expected_cid * 10); + try { + auto accessor = slot3->PinCells({target_uid}).get(); + if (!accessor) { + ADD_FAILURE() + << "T" << tid << " Op" << j + << ": PinCells returned null accessor for UID " + << target_uid; + test_failed = true; + break; + } + + TestCell* cell = accessor->get_cell_of(target_uid); + if (!cell) { + ADD_FAILURE() + << "T" << tid << " Op" << j + << ": get_cell_of returned null for UID " << target_uid; + test_failed = true; + break; + } + + if (cell->cid != expected_cid) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Incorrect CID for UID " << target_uid + << ". Slot: 3" + << ", Expected: " << expected_cid + << ", Got: " << cell->cid; + test_failed = true; + break; + } + + if (cell->data != expected_data) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Incorrect Data for UID " << target_uid + << ". Slot: 3" + << ", Expected: " << expected_data + << ", Got: " << cell->data; + test_failed = true; + break; + } + + if (ops_since_recreate >= recreate_interval) { + // std::cout << "Destroying SSSSSSslot3 at " << (void*)slot3.get() << std::endl; + slot3 = nullptr; + int sleep_ms = recreate_sleep_dist(gen); + std::this_thread::sleep_for( + std::chrono::milliseconds(sleep_ms)); + slot3 = create_new_slot3(); + ops_since_recreate = 0; + } else { + ops_since_recreate++; + int sleep_ms = sleep_dist(gen); + std::this_thread::sleep_for( + std::chrono::milliseconds(sleep_ms)); + } + } catch (const std::exception& e) { + ADD_FAILURE() + << "T" << tid << " Op" << j << ": Exception for UID " + << target_uid << ", Slot: 3" + << ". What: " << e.what(); + test_failed = true; + } catch (...) { + ADD_FAILURE() << "T" << tid << " Op" << j + << ": Unknown exception for UID " << target_uid + << ", Slot: 3"; + test_failed = true; + } + } + })); + + // 4. Wait for all threads to complete + try { + folly::collectAll(futures).get(); + } catch (const std::exception& e) { + FAIL() << "Exception waiting for thread pool completion: " << e.what(); + } catch (...) { + FAIL() << "Unknown exception waiting for thread pool completion."; + } + + ASSERT_FALSE(test_failed.load()) + << "Test failed due to assertion failures within threads."; + + ResourceUsage final_memory_usage = + DListTestFriend::get_used_memory(*dlist_); + + // bonus cell may cause memory usage to exceed the limit + if (!with_bonus_cells) { + EXPECT_LE(final_memory_usage.memory_bytes, NEW_LIMIT.memory_bytes) + << "Final memory usage (" << final_memory_usage.memory_bytes + << ") exceeds the limit (" << NEW_LIMIT.memory_bytes + << ") after concurrent access."; + } + + DListTestFriend::verify_integrity(dlist_.get()); +} + +INSTANTIATE_TEST_SUITE_P(BonusCellParam, + CacheSlotConcurrentTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "WithBonusCells" + : "NoBonusCells"; + }); diff --git a/internal/core/unittest/test_cachinglayer/test_dlist.cpp b/internal/core/unittest/test_cachinglayer/test_dlist.cpp new file mode 100644 index 0000000000..8a58240e5c --- /dev/null +++ b/internal/core/unittest/test_cachinglayer/test_dlist.cpp @@ -0,0 +1,610 @@ +#include +#include + +#include +#include +#include + +#include "cachinglayer/lrucache/DList.h" +#include "cachinglayer/Utils.h" +#include "mock_list_node.h" +#include "cachinglayer_test_utils.h" + +using namespace milvus::cachinglayer; +using namespace milvus::cachinglayer::internal; +using ::testing::StrictMock; +using DLF = DListTestFriend; + +class DListTest : public ::testing::Test { + protected: + ResourceUsage initial_limit{100, 50}; + DList::TouchConfig touch_config{{std::chrono::seconds(1)}}; + std::unique_ptr dlist; + // Keep track of nodes to prevent them from being deleted prematurely + std::vector> managed_nodes; + + void + SetUp() override { + dlist = std::make_unique(initial_limit, touch_config); + managed_nodes.clear(); + } + + void + TearDown() override { + managed_nodes.clear(); + dlist.reset(); + } + + // Helper to create a mock node, simulate loading it, and add it to the list. + // Returns a raw pointer, but ownership is managed by the shared_ptr in managed_nodes. + MockListNode* + add_and_load_node(ResourceUsage size, + const std::string& key = "key", + cid_t cid = 0, + int pin_count = 0) { + // Check if adding this node would exceed capacity before creating/adding it. + // We want to use add_and_load_node to create a DList in valid state. + ResourceUsage current_usage = get_used_memory(); + ResourceUsage limit = DLF::get_max_memory(*dlist); + if (!limit.CanHold(current_usage + size)) { + throw std::invalid_argument( + "Adding this node would exceed capacity"); + } + + auto node_ptr = std::make_shared>( + dlist.get(), size, key, cid); + managed_nodes.push_back(node_ptr); + MockListNode* node = node_ptr.get(); + + node->test_set_state(ListNode::State::LOADED); + node->test_set_pin_count(pin_count); + + // Manually adjust used memory and list pointers + DLF::test_add_used_memory(dlist.get(), size); + DLF::test_push_head(dlist.get(), node); + + return node; + } + + ResourceUsage + get_used_memory() const { + return DLF::get_used_memory(*dlist); + } + + // void + // DLF::verify_list(dlist.get(), std::vector nodes) const { + // EXPECT_EQ(nodes.front(), DLF::get_tail(*dlist)); + // EXPECT_EQ(nodes.back(), DLF::get_head(*dlist)); + // for (size_t i = 0; i < nodes.size() - 1; ++i) { + // auto current = nodes[i]; + // auto expected_prev = i == 0 ? nullptr : nodes[i - 1]; + // auto expected_next = i == nodes.size() - 1 ? nullptr : nodes[i + 1]; + // EXPECT_EQ(current->test_get_prev(), expected_prev); + // EXPECT_EQ(current->test_get_next(), expected_next); + // } + // } +}; + +TEST_F(DListTest, Initialization) { + EXPECT_TRUE(dlist->IsEmpty()); + EXPECT_EQ(get_used_memory(), ResourceUsage{}); + EXPECT_EQ(DLF::get_head(*dlist), nullptr); + EXPECT_EQ(DLF::get_tail(*dlist), nullptr); +} + +TEST_F(DListTest, UpdateLimitIncrease) { + MockListNode* node1 = add_and_load_node({10, 5}); + EXPECT_EQ(get_used_memory(), node1->size()); + + ResourceUsage new_limit{200, 100}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), node1->size()); + DLF::verify_list(dlist.get(), {node1}); +} + +TEST_F(DListTest, UpdateLimitDecreaseNoEviction) { + MockListNode* node1 = add_and_load_node({10, 5}); + ResourceUsage current_usage = node1->size(); + ASSERT_EQ(get_used_memory(), current_usage); + + ResourceUsage new_limit{50, 25}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), current_usage); + DLF::verify_list(dlist.get(), {node1}); +} + +TEST_F(DListTest, UpdateLimitDecreaseWithEvictionLRU) { + MockListNode* node1 = add_and_load_node({50, 20}, "key1"); + MockListNode* node2 = add_and_load_node({50, 30}, "key2"); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + EXPECT_EQ(get_used_memory(), usage_node1 + usage_node2); + EXPECT_EQ(get_used_memory(), DLF::get_max_memory(*dlist)); + + // Expect node1 to be evicted + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(0); + + ResourceUsage new_limit{70, 40}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), usage_node2); + DLF::verify_list(dlist.get(), {node2}); + EXPECT_FALSE(dlist->IsEmpty()); +} + +TEST_F(DListTest, UpdateLimitDecreaseWithEvictionMultiple) { + MockListNode* node1 = add_and_load_node({30, 10}, "key1"); + MockListNode* node2 = add_and_load_node({30, 10}, "key2"); + MockListNode* node3 = add_and_load_node({30, 10}, "key3"); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + ResourceUsage usage_node3 = node3->size(); + DLF::verify_list(dlist.get(), {node1, node2, node3}); + ASSERT_EQ(get_used_memory(), usage_node1 + usage_node2 + usage_node3); + + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(1); + EXPECT_CALL(*node3, clear_data()).Times(0); + + ResourceUsage new_limit{40, 15}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), usage_node3); + DLF::verify_list(dlist.get(), {node3}); +} + +TEST_F(DListTest, UpdateLimitSkipsPinned) { + MockListNode* node1 = add_and_load_node({40, 15}, "key1", 0, 1); + MockListNode* node2 = add_and_load_node({50, 25}, "key2"); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(get_used_memory(), usage_node1 + usage_node2); + + EXPECT_CALL(*node1, clear_data()).Times(0); + EXPECT_CALL(*node2, clear_data()).Times(1); + + ResourceUsage new_limit{70, 40}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), usage_node1); + DLF::verify_list(dlist.get(), {node1}); +} + +TEST_F(DListTest, UpdateLimitToZero) { + MockListNode* node1 = add_and_load_node({10, 0}); + MockListNode* node2 = add_and_load_node({0, 5}); + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(1); + + EXPECT_TRUE(dlist->UpdateLimit({1, 1})); + + EXPECT_EQ(get_used_memory(), ResourceUsage{}); + EXPECT_TRUE(dlist->IsEmpty()); +} + +TEST_F(DListTest, UpdateLimitInvalid) { + EXPECT_THROW(dlist->UpdateLimit({-10, 0}), std::invalid_argument); + EXPECT_THROW(dlist->UpdateLimit({0, -5}), std::invalid_argument); +} + +TEST_F(DListTest, ReserveMemorySufficient) { + ResourceUsage size{20, 10}; + EXPECT_TRUE(dlist->reserveMemory(size)); + EXPECT_EQ(get_used_memory(), size); +} + +TEST_F(DListTest, ReserveMemoryRequiresEviction) { + MockListNode* node1 = add_and_load_node({40, 15}, "key1"); + MockListNode* node2 = add_and_load_node({50, 25}, "key2"); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + + ASSERT_EQ(get_used_memory(), usage_node1 + usage_node2); + + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(0); + + ResourceUsage reserve_size{20, 20}; + // Current used: 90, 40. Limit: 100, 50. Reserve: 20, 20. + // Potential total: 110, 60. Need to free >= 10 mem, 10 disk. + // Evicting node1 ({40, 15}) is sufficient. + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_node2 + reserve_size); + DLF::verify_list(dlist.get(), {node2}); +} + +TEST_F(DListTest, ReserveMemoryEvictPinnedSkipped) { + MockListNode* node_pinned = add_and_load_node({40, 15}, "key_pinned", 0, 1); + MockListNode* node_evict = add_and_load_node({50, 25}, "key_evict"); + ResourceUsage usage_pinned = node_pinned->size(); + ResourceUsage usage_evict = node_evict->size(); + DLF::verify_list(dlist.get(), {node_pinned, node_evict}); + + ASSERT_EQ(get_used_memory(), usage_pinned + usage_evict); + + EXPECT_CALL(*node_pinned, clear_data()).Times(0); + EXPECT_CALL(*node_evict, clear_data()).Times(1); + + ResourceUsage reserve_size{20, 20}; + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_pinned + reserve_size); + DLF::verify_list(dlist.get(), {node_pinned}); +} + +TEST_F(DListTest, ReserveMemoryEvictLockedSkipped) { + MockListNode* node_locked = add_and_load_node({40, 15}, "key_locked"); + MockListNode* node_evict = add_and_load_node({50, 25}, "key_evict"); + ResourceUsage usage_locked = node_locked->size(); + ResourceUsage usage_evict = node_evict->size(); + DLF::verify_list(dlist.get(), {node_locked, node_evict}); + + ASSERT_EQ(get_used_memory(), usage_locked + usage_evict); + + EXPECT_CALL(*node_locked, clear_data()).Times(0); + EXPECT_CALL(*node_evict, clear_data()).Times(1); + + // Simulate locking the node during eviction attempt + std::unique_lock locked_node_lock(node_locked->test_get_mutex()); + + ResourceUsage reserve_size{20, 20}; + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + locked_node_lock.unlock(); + + EXPECT_EQ(get_used_memory(), usage_locked + reserve_size); + DLF::verify_list(dlist.get(), {node_locked}); +} + +TEST_F(DListTest, ReserveMemoryInsufficientEvenWithEviction) { + MockListNode* node1 = add_and_load_node({10, 5}); + ResourceUsage usage_node1 = node1->size(); + ASSERT_EQ(get_used_memory(), usage_node1); + + ResourceUsage reserve_size{200, 100}; + + EXPECT_FALSE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_node1); + EXPECT_FALSE(dlist->IsEmpty()); +} + +TEST_F(DListTest, TouchItemMovesToHead) { + MockListNode* node1 = add_and_load_node({10, 0}, "key1"); + MockListNode* node2 = add_and_load_node({10, 0}, "key2"); + MockListNode* node3 = add_and_load_node({10, 0}, "key3"); + + DLF::verify_list(dlist.get(), {node1, node2, node3}); + + { + std::unique_lock node_lock(node1->test_get_mutex()); + dlist->touchItem(node1); + } + + DLF::verify_list(dlist.get(), {node2, node3, node1}); +} + +TEST_F(DListTest, TouchItemRefreshWindow) { + MockListNode* node1 = add_and_load_node({10, 0}, "key1"); + MockListNode* node2 = add_and_load_node({10, 0}, "key2"); + + DLF::verify_list(dlist.get(), {node1, node2}); + + { + std::unique_lock node_lock(node1->test_get_mutex()); + dlist->touchItem(node1); + } + DLF::verify_list(dlist.get(), {node2, node1}); + + { + std::unique_lock node_lock(node1->test_get_mutex()); + dlist->touchItem(node1); + } + DLF::verify_list(dlist.get(), {node2, node1}); + + std::this_thread::sleep_for(touch_config.refresh_window + + std::chrono::milliseconds(100)); + + { + std::unique_lock node_lock(node1->test_get_mutex()); + dlist->touchItem(node1); + } + DLF::verify_list(dlist.get(), {node2, node1}); + + std::this_thread::sleep_for(touch_config.refresh_window + + std::chrono::milliseconds(100)); + + { + std::unique_lock node_lock(node2->test_get_mutex()); + dlist->touchItem(node2); + } + DLF::verify_list(dlist.get(), {node1, node2}); +} + +TEST_F(DListTest, releaseMemory) { + ResourceUsage initial_size{30, 15}; + DLF::test_add_used_memory(dlist.get(), initial_size); + ASSERT_EQ(get_used_memory(), initial_size); + + ResourceUsage failed_load_size{10, 5}; + dlist->releaseMemory(failed_load_size); + + EXPECT_EQ(get_used_memory(), initial_size - failed_load_size); +} + +TEST_F(DListTest, ReserveMemoryEvictOnlyMemoryNeeded) { + initial_limit = {100, 100}; + EXPECT_TRUE(dlist->UpdateLimit(initial_limit)); + + MockListNode* node_disk_only = add_and_load_node({0, 50}, "disk_only"); + MockListNode* node_mixed = add_and_load_node({50, 50}, "mixed"); + ResourceUsage usage_disk = node_disk_only->size(); + ResourceUsage usage_mixed = node_mixed->size(); + DLF::verify_list(dlist.get(), {node_disk_only, node_mixed}); + ASSERT_EQ(get_used_memory(), usage_disk + usage_mixed); + + EXPECT_CALL(*node_disk_only, clear_data()).Times(0); + EXPECT_CALL(*node_mixed, clear_data()).Times(1); + + // node_disk_only is at tail, but it contains no memory, thus evicting it does not help. + // We need to evict node_mixed to free up memory. + ResourceUsage reserve_size{60, 0}; + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_disk + reserve_size); + DLF::verify_list(dlist.get(), {node_disk_only}); +} + +TEST_F(DListTest, ReserveMemoryEvictOnlyDiskNeeded) { + initial_limit = {100, 100}; + EXPECT_TRUE(dlist->UpdateLimit(initial_limit)); + + MockListNode* node_mem_only = add_and_load_node({50, 0}, "mem_only"); + MockListNode* node_mixed = add_and_load_node({50, 50}, "mixed"); + ResourceUsage usage_mem = node_mem_only->size(); + ResourceUsage usage_mixed = node_mixed->size(); + DLF::verify_list(dlist.get(), {node_mem_only, node_mixed}); + ASSERT_EQ(get_used_memory(), usage_mem + usage_mixed); + + EXPECT_CALL(*node_mem_only, clear_data()).Times(0); + EXPECT_CALL(*node_mixed, clear_data()).Times(1); + + ResourceUsage reserve_size{0, 60}; + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_mem + reserve_size); + DLF::verify_list(dlist.get(), {node_mem_only}); +} + +TEST_F(DListTest, ReserveMemoryEvictBothNeeded) { + initial_limit = {100, 100}; + EXPECT_TRUE(dlist->UpdateLimit(initial_limit)); + + MockListNode* node1 = add_and_load_node({30, 10}, "node1"); + MockListNode* node2 = add_and_load_node({10, 30}, "node2"); + MockListNode* node3 = add_and_load_node({50, 50}, "node3"); + ResourceUsage usage1 = node1->size(); + ResourceUsage usage2 = node2->size(); + ResourceUsage usage3 = node3->size(); + DLF::verify_list(dlist.get(), {node1, node2, node3}); + ASSERT_EQ(get_used_memory(), usage1 + usage2 + usage3); + + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(1); + EXPECT_CALL(*node3, clear_data()).Times(0); + + ResourceUsage reserve_size{50, 50}; + EXPECT_TRUE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage3 + reserve_size); + DLF::verify_list(dlist.get(), {node3}); +} + +TEST_F(DListTest, ReserveMemoryFailsAllPinned) { + MockListNode* node1 = add_and_load_node({40, 15}, "key1", 0, 1); + MockListNode* node2 = add_and_load_node({50, 25}, "key2", 0, 1); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(get_used_memory(), usage_node1 + usage_node2); + + EXPECT_CALL(*node1, clear_data()).Times(0); + EXPECT_CALL(*node2, clear_data()).Times(0); + + ResourceUsage reserve_size{20, 20}; + EXPECT_FALSE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_node1 + usage_node2); + DLF::verify_list(dlist.get(), {node1, node2}); +} + +TEST_F(DListTest, ReserveMemoryFailsAllLocked) { + MockListNode* node1 = add_and_load_node({40, 15}, "key1"); + MockListNode* node2 = add_and_load_node({50, 25}, "key2"); + ResourceUsage usage_node1 = node1->size(); + ResourceUsage usage_node2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(get_used_memory(), usage_node1 + usage_node2); + + std::unique_lock lock1(node1->test_get_mutex()); + std::unique_lock lock2(node2->test_get_mutex()); + + EXPECT_CALL(*node1, clear_data()).Times(0); + EXPECT_CALL(*node2, clear_data()).Times(0); + + ResourceUsage reserve_size{20, 20}; + EXPECT_FALSE(dlist->reserveMemory(reserve_size)); + + lock1.unlock(); + lock2.unlock(); + + EXPECT_EQ(get_used_memory(), usage_node1 + usage_node2); + DLF::verify_list(dlist.get(), {node1, node2}); +} + +TEST_F(DListTest, ReserveMemoryFailsSpecificPinned) { + MockListNode* node_evict = + add_and_load_node({80, 40}, "evict_candidate", 0, 1); + MockListNode* node_small = add_and_load_node({10, 5}, "small"); + ResourceUsage usage_evict = node_evict->size(); + ResourceUsage usage_small = node_small->size(); + DLF::verify_list(dlist.get(), {node_evict, node_small}); + ASSERT_EQ(get_used_memory(), usage_evict + usage_small); + + EXPECT_CALL(*node_evict, clear_data()).Times(0); + EXPECT_CALL(*node_small, clear_data()).Times(0); + + ResourceUsage reserve_size{20, 20}; + EXPECT_FALSE(dlist->reserveMemory(reserve_size)); + + EXPECT_EQ(get_used_memory(), usage_evict + usage_small); + DLF::verify_list(dlist.get(), {node_evict, node_small}); +} + +TEST_F(DListTest, ReserveMemoryFailsSpecificLocked) { + MockListNode* node_evict = add_and_load_node({80, 40}, "evict_candidate"); + MockListNode* node_small = add_and_load_node({10, 5}, "small"); + ResourceUsage usage_evict = node_evict->size(); + ResourceUsage usage_small = node_small->size(); + DLF::verify_list(dlist.get(), {node_evict, node_small}); + ASSERT_EQ(get_used_memory(), usage_evict + usage_small); + + std::unique_lock lock_evict(node_evict->test_get_mutex()); + + EXPECT_CALL(*node_evict, clear_data()).Times(0); + EXPECT_CALL(*node_small, clear_data()).Times(0); + + ResourceUsage reserve_size{20, 20}; + EXPECT_FALSE(dlist->reserveMemory(reserve_size)); + + lock_evict.unlock(); + + EXPECT_EQ(get_used_memory(), usage_evict + usage_small); + DLF::verify_list(dlist.get(), {node_evict, node_small}); +} + +TEST_F(DListTest, TouchItemHeadOutsideWindow) { + MockListNode* node1 = add_and_load_node({10, 0}, "key1"); + MockListNode* node2 = add_and_load_node({10, 0}, "key2"); + DLF::verify_list(dlist.get(), {node1, node2}); + + std::this_thread::sleep_for(touch_config.refresh_window + + std::chrono::milliseconds(100)); + + { + std::unique_lock node_lock(node2->test_get_mutex()); + dlist->touchItem(node2); + } + + DLF::verify_list(dlist.get(), {node1, node2}); +} + +TEST_F(DListTest, RemoveItemFromList) { + MockListNode* node1 = add_and_load_node({10, 0}, "key1"); + MockListNode* node2 = add_and_load_node({10, 0}, "key2"); + DLF::verify_list(dlist.get(), {node1, node2}); + + { + std::unique_lock node_lock(node1->test_get_mutex()); + dlist->removeItem(node1, node1->size()); + } + + DLF::verify_list(dlist.get(), {node2}); + EXPECT_EQ(get_used_memory(), node2->size()); +} + +TEST_F(DListTest, PopItemNotPresent) { + MockListNode* node1 = add_and_load_node({10, 0}, "key1"); + MockListNode* node2 = add_and_load_node({10, 0}, "key2"); + ResourceUsage initial_usage = get_used_memory(); + DLF::verify_list(dlist.get(), {node1, node2}); + + auto orphan_node_ptr = std::make_unique>( + dlist.get(), ResourceUsage{10, 0}, "orphan", 0); + MockListNode* orphan_node = orphan_node_ptr.get(); + + { + std::unique_lock node_lock(orphan_node->test_get_mutex()); + EXPECT_NO_THROW(DLF::test_pop_item(dlist.get(), orphan_node)); + } + + DLF::verify_list(dlist.get(), {node1, node2}); + EXPECT_EQ(get_used_memory(), initial_usage); +} + +TEST_F(DListTest, UpdateLimitIncreaseMemDecreaseDisk) { + MockListNode* node1 = add_and_load_node({20, 30}, "node1"); + MockListNode* node2 = add_and_load_node({30, 10}, "node2"); + ResourceUsage usage1 = node1->size(); + ResourceUsage usage2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(get_used_memory(), usage1 + usage2); + + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(0); + + ResourceUsage new_limit{200, 35}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + + EXPECT_EQ(get_used_memory(), usage2); + DLF::verify_list(dlist.get(), {node2}); + EXPECT_EQ(DLF::get_max_memory(*dlist), new_limit); +} + +TEST_F(DListTest, EvictedNodeDestroyed) { + MockListNode* node1 = add_and_load_node({40, 15}, "node1"); + MockListNode* node2 = add_and_load_node({50, 25}, "node2"); + ResourceUsage usage1 = node1->size(); + ResourceUsage usage2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(managed_nodes.size(), 2); + ASSERT_EQ(get_used_memory(), usage1 + usage2); + + EXPECT_CALL(*node1, clear_data()).Times(1); + EXPECT_CALL(*node2, clear_data()).Times(0); + ResourceUsage new_limit{70, 40}; + EXPECT_TRUE(dlist->UpdateLimit(new_limit)); + DLF::verify_list(dlist.get(), {node2}); + ResourceUsage memory_after_eviction = get_used_memory(); + ASSERT_EQ(memory_after_eviction, usage2); + + // destroy node1 by removing its shared_ptr + // node1's destructor should not decrement used_memory_ again + auto it = std::find_if(managed_nodes.begin(), + managed_nodes.end(), + [&](const auto& ptr) { return ptr.get() == node1; }); + ASSERT_NE(it, managed_nodes.end()); + managed_nodes.erase(it); + + EXPECT_EQ(get_used_memory(), memory_after_eviction); + DLF::verify_list(dlist.get(), {node2}); +} + +TEST_F(DListTest, NodeInListDestroyed) { + MockListNode* node1 = add_and_load_node({40, 15}, "node1"); + MockListNode* node2 = add_and_load_node({50, 25}, "node2"); + ResourceUsage usage1 = node1->size(); + ResourceUsage usage2 = node2->size(); + DLF::verify_list(dlist.get(), {node1, node2}); + ASSERT_EQ(managed_nodes.size(), 2); + ResourceUsage memory_before_destroy = get_used_memory(); + ASSERT_EQ(memory_before_destroy, usage1 + usage2); + + // destroy node1 by removing its shared_ptr + // node1's destructor should decrement used_memory_ by node1->size() and remove node1 from the list + auto it = std::find_if(managed_nodes.begin(), + managed_nodes.end(), + [&](const auto& ptr) { return ptr.get() == node1; }); + ASSERT_NE(it, managed_nodes.end()); + managed_nodes.erase(it); + + EXPECT_EQ(get_used_memory(), memory_before_destroy - usage1); + DLF::verify_list(dlist.get(), {node2}); +} diff --git a/internal/core/unittest/test_chunk.cpp b/internal/core/unittest/test_chunk.cpp index 8d8ee7f44f..8b3042ecfa 100644 --- a/internal/core/unittest/test_chunk.cpp +++ b/internal/core/unittest/test_chunk.cpp @@ -67,7 +67,8 @@ TEST(chunk, test_int64_field) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto span = std::dynamic_pointer_cast(chunk)->Span(); + auto fixed_chunk = static_cast(chunk.get()); + auto span = fixed_chunk->Span(); EXPECT_EQ(span.row_count(), data.size()); for (size_t i = 0; i < data.size(); ++i) { auto n = *(int64_t*)((char*)span.data() + i * span.element_sizeof()); @@ -109,8 +110,8 @@ TEST(chunk, test_variable_field) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto views = std::dynamic_pointer_cast(chunk)->StringViews( - std::nullopt); + auto string_chunk = static_cast(chunk.get()); + auto views = string_chunk->StringViews(std::nullopt); for (size_t i = 0; i < data.size(); ++i) { EXPECT_EQ(views.first[i], data[i]); } @@ -154,8 +155,8 @@ TEST(chunk, test_variable_field_nullable) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto views = std::dynamic_pointer_cast(chunk)->StringViews( - std::nullopt); + auto string_chunk = static_cast(chunk.get()); + auto views = string_chunk->StringViews(std::nullopt); for (size_t i = 0; i < data.size(); ++i) { EXPECT_EQ(views.second[i], validity[i]); if (validity[i]) { @@ -211,10 +212,9 @@ TEST(chunk, test_json_field) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); + auto json_chunk = static_cast(chunk.get()); { - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->StringViews( - std::nullopt); + auto [views, valid] = json_chunk->StringViews(std::nullopt); EXPECT_EQ(row_num, views.size()); for (size_t i = 0; i < row_num; ++i) { EXPECT_EQ(views[i], data[i].data()); @@ -225,8 +225,7 @@ TEST(chunk, test_json_field) { auto start = 10; auto len = 20; auto [views, valid] = - std::dynamic_pointer_cast(chunk)->StringViews( - std::make_pair(start, len)); + json_chunk->StringViews(std::make_pair(start, len)); EXPECT_EQ(len, views.size()); for (size_t i = 0; i < len; ++i) { EXPECT_EQ(views[i], data[i].data()); @@ -243,10 +242,9 @@ TEST(chunk, test_json_field) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); + auto json_chunk = static_cast(chunk.get()); { - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->StringViews( - std::nullopt); + auto [views, valid] = json_chunk->StringViews(std::nullopt); EXPECT_EQ(row_num, views.size()); for (size_t i = 0; i < row_num; ++i) { EXPECT_EQ(views[i], data[i].data()); @@ -257,8 +255,7 @@ TEST(chunk, test_json_field) { auto start = 10; auto len = 20; auto [views, valid] = - std::dynamic_pointer_cast(chunk)->StringViews( - std::make_pair(start, len)); + json_chunk->StringViews(std::make_pair(start, len)); EXPECT_EQ(len, views.size()); for (size_t i = 0; i < len; ++i) { EXPECT_EQ(views[i], data[i].data()); @@ -268,26 +265,20 @@ TEST(chunk, test_json_field) { { auto start = -1; auto len = 5; - EXPECT_THROW( - std::dynamic_pointer_cast(chunk)->StringViews( - std::make_pair(start, len)), - milvus::SegcoreError); + EXPECT_THROW(json_chunk->StringViews(std::make_pair(start, len)), + milvus::SegcoreError); } { auto start = 0; auto len = row_num + 1; - EXPECT_THROW( - std::dynamic_pointer_cast(chunk)->StringViews( - std::make_pair(start, len)), - milvus::SegcoreError); + EXPECT_THROW(json_chunk->StringViews(std::make_pair(start, len)), + milvus::SegcoreError); } { auto start = 95; auto len = 11; - EXPECT_THROW( - std::dynamic_pointer_cast(chunk)->StringViews( - std::make_pair(start, len)), - milvus::SegcoreError); + EXPECT_THROW(json_chunk->StringViews(std::make_pair(start, len)), + milvus::SegcoreError); } } } @@ -329,7 +320,7 @@ TEST(chunk, test_null_int64) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto fixed_chunk = std::dynamic_pointer_cast(chunk); + auto fixed_chunk = static_cast(chunk.get()); auto span = fixed_chunk->Span(); EXPECT_EQ(span.row_count(), data.size()); @@ -390,8 +381,8 @@ TEST(chunk, test_array) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->Views(std::nullopt); + auto array_chunk = static_cast(chunk.get()); + auto [views, valid] = array_chunk->Views(std::nullopt); EXPECT_EQ(views.size(), 1); auto& arr = views[0]; for (size_t i = 0; i < arr.length(); ++i) { @@ -453,8 +444,8 @@ TEST(chunk, test_null_array) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->Views(std::nullopt); + auto array_chunk = static_cast(chunk.get()); + auto [views, valid] = array_chunk->Views(std::nullopt); EXPECT_EQ(views.size(), array_count); EXPECT_EQ(valid.size(), array_count); @@ -527,10 +518,9 @@ TEST(chunk, test_array_views) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, 1, array_vec); - + auto array_chunk = static_cast(chunk.get()); { - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->Views(std::nullopt); + auto [views, valid] = array_chunk->Views(std::nullopt); EXPECT_EQ(views.size(), array_count); for (auto i = 0; i < array_count; i++) { auto& arr = views[i]; @@ -543,9 +533,7 @@ TEST(chunk, test_array_views) { { auto start = 2; auto len = 5; - auto [views, valid] = - std::dynamic_pointer_cast(chunk)->Views( - std::make_pair(start, len)); + auto [views, valid] = array_chunk->Views(std::make_pair(start, len)); EXPECT_EQ(views.size(), len); for (auto i = 0; i < len; i++) { auto& arr = views[i]; @@ -558,22 +546,19 @@ TEST(chunk, test_array_views) { { auto start = -1; auto len = 5; - EXPECT_THROW(std::dynamic_pointer_cast(chunk)->Views( - std::make_pair(start, len)), + EXPECT_THROW(array_chunk->Views(std::make_pair(start, len)), milvus::SegcoreError); } { auto start = 0; auto len = array_count + 1; - EXPECT_THROW(std::dynamic_pointer_cast(chunk)->Views( - std::make_pair(start, len)), + EXPECT_THROW(array_chunk->Views(std::make_pair(start, len)), milvus::SegcoreError); } { auto start = 5; auto len = 7; - EXPECT_THROW(std::dynamic_pointer_cast(chunk)->Views( - std::make_pair(start, len)), + EXPECT_THROW(array_chunk->Views(std::make_pair(start, len)), milvus::SegcoreError); } } @@ -615,7 +600,8 @@ TEST(chunk, test_sparse_float) { std::nullopt); arrow::ArrayVector array_vec = read_single_column_batches(rb_reader); auto chunk = create_chunk(field_meta, kTestSparseDim, array_vec); - auto vec = std::dynamic_pointer_cast(chunk)->Vec(); + auto vec_chunk = static_cast(chunk.get()); + auto vec = vec_chunk->Vec(); for (size_t i = 0; i < n_rows; ++i) { auto v1 = vec[i]; auto v2 = vecs[i]; diff --git a/internal/core/unittest/test_chunk_cache.cpp b/internal/core/unittest/test_chunk_cache.cpp deleted file mode 100644 index 8b6150d1fc..0000000000 --- a/internal/core/unittest/test_chunk_cache.cpp +++ /dev/null @@ -1,251 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include - -#include -#include -#include -#include -#include - -#include "common/Consts.h" -#include "common/FieldMeta.h" -#include "common/Types.h" -#include "fmt/format.h" -#include "common/Schema.h" -#include "gtest/gtest.h" -#include "knowhere/sparse_utils.h" -#include "mmap/Column.h" -#include "test_utils/DataGen.h" -#include "test_utils/storage_test_utils.h" -#include "storage/ChunkCache.h" -#include "storage/LocalChunkManagerSingleton.h" - -#define DEFAULT_READ_AHEAD_POLICY "willneed" -class ChunkCacheTest : public testing::TestWithParam { - protected: - void - SetUp() override { - mcm = milvus::storage::MmapManager::GetInstance().GetMmapChunkManager(); - mcm->Register(descriptor); - - N = 10000; - dim = 128; - auto dense_metric_type = knowhere::metric::L2; - auto sparse_metric_type = knowhere::metric::IP; - - auto schema = std::make_shared(); - auto fake_dense_vec_id = schema->AddDebugField( - "fakevec", milvus::DataType::VECTOR_FLOAT, dim, dense_metric_type); - auto i64_fid = - schema->AddDebugField("counter", milvus::DataType::INT64); - auto fake_sparse_vec_id = - schema->AddDebugField("fakevec_sparse", - milvus::DataType::VECTOR_SPARSE_FLOAT, - dim, - sparse_metric_type); - schema->set_primary_field_id(i64_fid); - - auto dataset = milvus::segcore::DataGen(schema, N); - - auto dense_field_data_meta = - milvus::storage::FieldDataMeta{1, 2, 3, fake_dense_vec_id.get()}; - auto sparse_field_data_meta = - milvus::storage::FieldDataMeta{1, 2, 3, fake_sparse_vec_id.get()}; - dense_field_meta = milvus::FieldMeta(milvus::FieldName("fakevec"), - fake_dense_vec_id, - milvus::DataType::VECTOR_FLOAT, - dim, - dense_metric_type, - false, - std::nullopt); - sparse_field_meta = - milvus::FieldMeta(milvus::FieldName("fakevec_sparse"), - fake_sparse_vec_id, - milvus::DataType::VECTOR_SPARSE_FLOAT, - dim, - sparse_metric_type, - false, - std::nullopt); - - lcm = milvus::storage::LocalChunkManagerSingleton::GetInstance() - .GetChunkManager(); - dense_data = dataset.get_col(fake_dense_vec_id); - sparse_data = dataset.get_col>( - fake_sparse_vec_id); - - auto data_slices = std::vector{dense_data.data()}; - auto slice_sizes = std::vector{static_cast(N)}; - auto slice_names = std::vector{dense_file_name}; - PutFieldData(lcm.get(), - data_slices, - slice_sizes, - slice_names, - dense_field_data_meta, - dense_field_meta); - - data_slices = std::vector{sparse_data.data()}; - slice_sizes = std::vector{static_cast(N)}; - slice_names = std::vector{sparse_file_name}; - PutFieldData(lcm.get(), - data_slices, - slice_sizes, - slice_names, - sparse_field_data_meta, - sparse_field_meta); - } - void - TearDown() override { - mcm->UnRegister(descriptor); - } - const char* dense_file_name = "chunk_cache_test/insert_log/2/101/1000000"; - const char* sparse_file_name = "chunk_cache_test/insert_log/2/102/1000000"; - milvus::storage::MmapChunkManagerPtr mcm; - milvus::segcore::SegcoreConfig config; - milvus::storage::MmapChunkDescriptorPtr descriptor = - std::shared_ptr( - new milvus::storage::MmapChunkDescriptor( - {101, SegmentType::Sealed})); - - int N; - int dim; - milvus::FieldMeta dense_field_meta = milvus::FieldMeta::RowIdMeta; - milvus::FixedVector dense_data; - milvus::FieldMeta sparse_field_meta = milvus::FieldMeta::RowIdMeta; - milvus::FixedVector> sparse_data; - std::shared_ptr lcm; -}; -INSTANTIATE_TEST_SUITE_P(ChunkCacheTestSuite, ChunkCacheTest, testing::Bool()); - -TEST_P(ChunkCacheTest, Read) { - auto cc = milvus::storage::MmapManager::GetInstance().GetChunkCache(); - - // validate dense data - std::shared_ptr dense_column; - - auto mmap_enabled = GetParam(); - - dense_column = cc->Read(dense_file_name, dense_field_meta, mmap_enabled); - - auto actual_dense = (const float*)(dense_column->Data(0)); - for (auto i = 0; i < N * dim; i++) { - AssertInfo(dense_data[i] == actual_dense[i], - fmt::format( - "expect {}, actual {}", dense_data[i], actual_dense[i])); - } - - // validate sparse data - std::shared_ptr sparse_column; - sparse_column = cc->Read(sparse_file_name, sparse_field_meta, mmap_enabled); - - auto expected_sparse_size = 0; - auto actual_sparse = - (const knowhere::sparse::SparseRow*)(sparse_column->Data(0)); - for (auto i = 0; i < N; i++) { - const auto& actual_sparse_row = actual_sparse[i]; - const auto& expect_sparse_row = sparse_data[i]; - AssertInfo( - actual_sparse_row.size() == expect_sparse_row.size(), - fmt::format("Incorrect size of sparse row: expect {}, actual {}", - expect_sparse_row.size(), - actual_sparse_row.size())); - auto bytes = actual_sparse_row.data_byte_size(); - AssertInfo( - memcmp(actual_sparse_row.data(), expect_sparse_row.data(), bytes) == - 0, - fmt::format("Incorrect data of sparse row: expect {}, actual {}", - expect_sparse_row.data(), - actual_sparse_row.data())); - expected_sparse_size += bytes; - } - - expected_sparse_size += (N + 7) / 8; - expected_sparse_size += sizeof(int64_t) * (N + 1); - - if (mmap_enabled) { - const uint32_t page_size = sysconf(_SC_PAGE_SIZE); - auto padding_size = (expected_sparse_size / page_size + - (expected_sparse_size % page_size != 0)) * - page_size - - expected_sparse_size; - expected_sparse_size += padding_size; - } - auto actual_sparse_size = sparse_column->DataByteSize(); - Assert(actual_sparse_size == expected_sparse_size); - - cc->Remove(dense_file_name); - cc->Remove(sparse_file_name); - lcm->Remove(dense_file_name); - lcm->Remove(sparse_file_name); -} - -TEST_P(ChunkCacheTest, TestMultithreads) { - auto cc = milvus::storage::MmapManager::GetInstance().GetChunkCache(); - - constexpr int threads = 16; - std::vector total_counts(threads); - auto mmap_enabled = GetParam(); - auto executor = [&](int thread_id) { - std::shared_ptr dense_column; - dense_column = - cc->Read(dense_file_name, dense_field_meta, mmap_enabled); - - auto actual_dense = (const float*)dense_column->Data(0); - for (auto i = 0; i < N * dim; i++) { - AssertInfo( - dense_data[i] == actual_dense[i], - fmt::format( - "expect {}, actual {}", dense_data[i], actual_dense[i])); - } - - std::shared_ptr sparse_column; - sparse_column = - cc->Read(sparse_file_name, sparse_field_meta, mmap_enabled); - - auto actual_sparse = - (const knowhere::sparse::SparseRow*)sparse_column->Data(0); - for (auto i = 0; i < N; i++) { - const auto& actual_sparse_row = actual_sparse[i]; - const auto& expect_sparse_row = sparse_data[i]; - AssertInfo(actual_sparse_row.size() == expect_sparse_row.size(), - fmt::format( - "Incorrect size of sparse row: expect {}, actual {}", - expect_sparse_row.size(), - actual_sparse_row.size())); - auto bytes = actual_sparse_row.data_byte_size(); - AssertInfo(memcmp(actual_sparse_row.data(), - expect_sparse_row.data(), - bytes) == 0, - fmt::format( - "Incorrect data of sparse row: expect {}, actual {}", - expect_sparse_row.data(), - actual_sparse_row.data())); - } - }; - std::vector pool; - for (int i = 0; i < threads; ++i) { - pool.emplace_back(executor, i); - } - for (auto& thread : pool) { - thread.join(); - } - - cc->Remove(dense_file_name); - cc->Remove(sparse_file_name); - lcm->Remove(dense_file_name); - lcm->Remove(sparse_file_name); -} diff --git a/internal/core/unittest/test_chunked_column.cpp b/internal/core/unittest/test_chunked_column.cpp index e17de1163f..a8b3a0fc7b 100644 --- a/internal/core/unittest/test_chunked_column.cpp +++ b/internal/core/unittest/test_chunked_column.cpp @@ -9,22 +9,33 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License +#include #include "common/Chunk.h" #include "gtest/gtest.h" #include "mmap/ChunkedColumn.h" +#include "test_cachinglayer/cachinglayer_test_utils.h" + namespace milvus { + TEST(test_chunked_column, test_get_chunkid) { - ChunkedColumn column; - std::vector chunk_row_nums = {10, 20, 30}; - for (auto row_num : chunk_row_nums) { + std::vector num_rows_per_chunk = {0, 10, 20, 30}; + auto num_chunks = num_rows_per_chunk.size(); + std::vector> chunks; + for (auto i = 0; i < num_chunks; ++i) { + auto row_num = num_rows_per_chunk[i]; auto chunk = - std::make_shared(row_num, 1, nullptr, 0, 4, false); - column.AddChunk(chunk); + std::make_unique(row_num, 1, nullptr, 0, 4, false); + chunks.push_back(std::move(chunk)); } + auto translator = std::make_unique( + num_rows_per_chunk, "test", std::move(chunks)); + FieldMeta field_meta( + FieldName("test"), FieldId(1), DataType::INT64, false, std::nullopt); + ChunkedColumn column(std::move(translator), field_meta); int offset = 0; - for (int i = 0; i < chunk_row_nums.size(); ++i) { - for (int j = 0; j < chunk_row_nums[i]; ++j) { + for (int i = 0; i < num_chunks; ++i) { + for (int j = 0; j < num_rows_per_chunk[i]; ++j) { auto [chunk_id, offset_in_chunk] = column.GetChunkIDByOffset(offset); ASSERT_EQ(chunk_id, i); @@ -33,4 +44,4 @@ TEST(test_chunked_column, test_get_chunkid) { } } } -} // namespace milvus \ No newline at end of file +} // namespace milvus diff --git a/internal/core/unittest/test_chunked_segment.cpp b/internal/core/unittest/test_chunked_segment.cpp index f64c72d591..544912b19e 100644 --- a/internal/core/unittest/test_chunked_segment.cpp +++ b/internal/core/unittest/test_chunked_segment.cpp @@ -11,13 +11,20 @@ #include #include + +#include +#include +#include +#include +#include +#include #include #include -#include "arrow/table_builder.h" + +#include #include "arrow/type_fwd.h" #include "common/BitsetView.h" #include "common/Consts.h" -#include "common/FieldDataInterface.h" #include "common/QueryInfo.h" #include "common/Schema.h" #include "common/Types.h" @@ -29,22 +36,18 @@ #include "knowhere/comp/index_param.h" #include "milvus-storage/common/constants.h" #include "mmap/ChunkedColumn.h" -#include "mmap/Types.h" #include "pb/plan.pb.h" #include "pb/schema.pb.h" #include "query/ExecPlanNodeVisitor.h" #include "query/SearchOnSealed.h" #include "segcore/SegcoreConfig.h" #include "segcore/SegmentSealed.h" +#include "storage/RemoteChunkManagerSingleton.h" #include "segcore/Types.h" #include "test_utils/DataGen.h" -#include -#include -#include -#include -#include -#include +#include "test_utils/storage_test_utils.h" +#include "test_cachinglayer/cachinglayer_test_utils.h" struct DeferRelease { using functype = std::function; @@ -72,12 +75,16 @@ TEST(test_chunk_segment, TestSearchOnSealed) { int total_row_count = chunk_num * chunk_size; int bitset_size = (total_row_count + 7) / 8; - auto column = std::make_shared(); auto schema = std::make_shared(); auto fakevec_id = schema->AddDebugField( "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::COSINE); + auto field_meta = schema->operator[](fakevec_id); + + std::vector> chunks; + std::vector num_rows_per_chunk; for (int i = 0; i < chunk_num; i++) { + num_rows_per_chunk.push_back(chunk_size); auto dataset = segcore::DataGen(schema, chunk_size); auto data = dataset.get_col(fakevec_id); auto buf_size = 4 * data.size(); @@ -86,11 +93,15 @@ TEST(test_chunk_segment, TestSearchOnSealed) { defer.AddDefer([buf]() { delete[] buf; }); memcpy(buf, data.data(), 4 * data.size()); - auto chunk = std::make_shared( - chunk_size, dim, buf, buf_size, 4, false); - column->AddChunk(chunk); + chunks.emplace_back(std::make_unique( + chunk_size, dim, buf, buf_size, 4, false)); } + auto translator = std::make_unique( + num_rows_per_chunk, "", std::move(chunks)); + auto column = + std::make_shared(std::move(translator), field_meta); + SearchInfo search_info; auto search_conf = knowhere::Json{ {knowhere::meta::METRIC_TYPE, knowhere::metric::COSINE}, @@ -112,15 +123,15 @@ TEST(test_chunk_segment, TestSearchOnSealed) { auto index_info = std::map{}; SearchResult search_result; - query::SearchOnSealed(*schema, - column, - search_info, - index_info, - query_data, - 1, - total_row_count, - bv, - search_result); + query::SearchOnSealedColumn(*schema, + column.get(), + search_info, + index_info, + query_data, + 1, + total_row_count, + bv, + search_result); std::set offsets; for (auto& offset : search_result.seg_offsets_) { @@ -137,15 +148,15 @@ TEST(test_chunk_segment, TestSearchOnSealed) { // test with group by search_info.group_by_field_id_ = fakevec_id; std::fill(bitset_data, bitset_data + bitset_size, 0); - query::SearchOnSealed(*schema, - column, - search_info, - index_info, - query_data, - 1, - total_row_count, - bv, - search_result); + query::SearchOnSealedColumn(*schema, + column.get(), + search_info, + index_info, + query_data, + 1, + total_row_count, + bv, + search_result); ASSERT_EQ(1, search_result.vector_iterators_->size()); @@ -172,7 +183,7 @@ class TestChunkSegment : public testing::TestWithParam { auto int64_fid = schema->AddDebugField("int64", DataType::INT64, true); auto pk_fid = schema->AddDebugField( - "pk", pk_is_string ? DataType::VARCHAR : DataType::INT64, true); + "pk", pk_is_string ? DataType::VARCHAR : DataType::INT64, false); auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR, true); auto str2_fid = @@ -201,7 +212,7 @@ class TestChunkSegment : public testing::TestWithParam { auto arrow_pk_field = arrow::field( "pk", pk_is_string ? arrow::utf8() : arrow::int64(), - true, + false, arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY}, {std::to_string(101)})); auto arrow_ts_field = arrow::field( @@ -245,20 +256,16 @@ class TestChunkSegment : public testing::TestWithParam { int start_id = 0; chunk_num = 2; - std::vector field_infos; - for (auto fid : field_ids) { - FieldDataInfo field_info; - field_info.field_id = fid.get(); - field_info.row_count = test_data_count * chunk_num; - field_infos.push_back(field_info); - } - std::vector str_data; for (int i = 0; i < test_data_count * chunk_num; i++) { str_data.push_back("test" + std::to_string(i)); } std::sort(str_data.begin(), str_data.end()); - std::vector validity(test_data_count, true); + + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + + std::unordered_map> field_data_map; // generate data for (int chunk_id = 0; chunk_id < chunk_num; @@ -266,50 +273,38 @@ class TestChunkSegment : public testing::TestWithParam { std::vector test_data(test_data_count); std::iota(test_data.begin(), test_data.end(), start_id); - auto builder = std::make_shared(); - auto status = builder->AppendValues( - test_data.begin(), test_data.end(), validity.begin()); - ASSERT_TRUE(status.ok()); - auto res = builder->Finish(); - ASSERT_TRUE(res.ok()); - std::shared_ptr arrow_int64; - arrow_int64 = res.ValueOrDie(); - - auto str_builder = std::make_shared(); - for (int i = 0; i < test_data_count; i++) { - auto status = str_builder->Append(str_data[start_id + i]); - ASSERT_TRUE(status.ok()); - } - std::shared_ptr arrow_str; - status = str_builder->Finish(&arrow_str); - ASSERT_TRUE(status.ok()); - for (int i = 0; i < arrow_fields.size(); i++) { - auto f = arrow_fields[i]; + // if i < 3, this is system field, thus must be int64. + // other fields include a pk field and 2 string fields. pk field is string if pk_is_string is true. + auto datatype = + i < 3 && (field_ids[i] != pk_fid || !pk_is_string) + ? DataType::INT64 + : DataType::VARCHAR; + FieldDataPtr field_data{nullptr}; + + if (datatype == DataType::INT64) { + field_data = + std::make_shared>(datatype, false); + field_data->FillFieldData(test_data.data(), + test_data_count); + } else { + field_data = std::make_shared>( + datatype, false); + field_data->FillFieldData(str_data.data() + start_id, + test_data_count); + } auto fid = field_ids[i]; - auto arrow_schema = - std::make_shared(arrow::FieldVector(1, f)); - - auto col = i < 3 && (field_ids[i] != pk_fid || !pk_is_string) - ? arrow_int64 - : arrow_str; - auto record_batch = arrow::RecordBatch::Make( - arrow_schema, arrow_int64->length(), {col}); - - auto res2 = arrow::RecordBatchReader::Make({record_batch}); - ASSERT_TRUE(res2.ok()); - auto arrow_reader = res2.ValueOrDie(); - - field_infos[i].arrow_reader_channel->push( - std::make_shared( - arrow_reader, nullptr, nullptr)); + field_data_map[fid].push_back(field_data); } } - - // load - for (int i = 0; i < field_infos.size(); i++) { - field_infos[i].arrow_reader_channel->close(); - segment->LoadFieldData(field_ids[i], field_infos[i]); + for (auto& [fid, field_datas] : field_data_map) { + auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + fid.get(), + field_datas, + cm); + segment->LoadFieldData(load_info); } } @@ -422,7 +417,8 @@ TEST_P(TestChunkSegment, TestCompareExpr) { create_index_info, file_manager_ctx); std::vector data(test_data_count * chunk_num); for (int i = 0; i < chunk_num; i++) { - auto d = segment->chunk_data(fid, i); + auto pw = segment->chunk_data(fid, i); + auto d = pw.get(); std::copy(d.data(), d.data() + test_data_count, data.begin() + i * test_data_count); diff --git a/internal/core/unittest/test_delete_record.cpp b/internal/core/unittest/test_delete_record.cpp index b01eeac7ab..621be46b99 100644 --- a/internal/core/unittest/test_delete_record.cpp +++ b/internal/core/unittest/test_delete_record.cpp @@ -12,20 +12,16 @@ #include #include -#include #include #include #include #include -#include #include -#include #include "segcore/DeletedRecord.h" -#include "segcore/SegmentGrowingImpl.h" - -#include "segcore/SegmentGrowingImpl.h" +#include "segcore/Record.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::segcore; @@ -34,21 +30,20 @@ TEST(DeleteMVCC, common_case) { auto schema = std::make_shared(); auto pk = schema->AddDebugField("pk", DataType::INT64); schema->set_primary_field_id(pk); - auto segment = CreateSealedSegment(schema); - ASSERT_EQ(0, segment->get_real_count()); // load insert: pk (0, 1, 2, 3, 4, 5, 6, 7, 8, 9) // with timestamp ts (0, 1, 2, 3, 4, 5, 6, 7, 8, 9) int64_t c = 10; auto dataset = DataGen(schema, c); auto pks = dataset.get_col(pk); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); ASSERT_EQ(c, segment->get_real_count()); - auto& insert_record = segment->get_insert_record(); + auto insert_record = segment->get_insert_record_slot(); + auto segment_ptr = segment.get(); DeletedRecord delete_record( - &insert_record, - [&insert_record](const PkType& pk, Timestamp timestamp) { - return insert_record.search_pk(pk, timestamp); + insert_record, + [segment_ptr](const PkType& pk, Timestamp timestamp) { + return segment_ptr->search_pk(pk, timestamp); }, 0); delete_record.set_sealed_row_count(c); @@ -160,7 +155,7 @@ TEST(DeleteMVCC, delete_exist_duplicate_pks) { schema->set_primary_field_id(i64_fid); auto N = 10; uint64_t seg_id = 101; - InsertRecord insert_record(*schema, N); + InsertRecord insert_record(*schema, N); DeletedRecord delete_record( &insert_record, [&insert_record](const PkType& pk, Timestamp timestamp) { @@ -275,7 +270,7 @@ TEST(DeleteMVCC, snapshot) { schema->set_primary_field_id(i64_fid); auto N = 50000; uint64_t seg_id = 101; - InsertRecord insert_record(*schema, N); + InsertRecord insert_record(*schema, N); DeletedRecord delete_record( &insert_record, [&insert_record](const PkType& pk, Timestamp timestamp) { @@ -323,7 +318,7 @@ TEST(DeleteMVCC, insert_after_snapshot) { schema->set_primary_field_id(i64_fid); auto N = 11000; uint64_t seg_id = 101; - InsertRecord insert_record(*schema, N); + InsertRecord insert_record(*schema, N); DeletedRecord delete_record( &insert_record, [&insert_record](const PkType& pk, Timestamp timestamp) { @@ -418,7 +413,7 @@ TEST(DeleteMVCC, perform) { schema->set_primary_field_id(i64_fid); auto N = 1000000; uint64_t seg_id = 101; - InsertRecord insert_record(*schema, N); + InsertRecord insert_record(*schema, N); DeletedRecord delete_record( &insert_record, [&insert_record](const PkType& pk, Timestamp timestamp) { diff --git a/internal/core/unittest/test_exec.cpp b/internal/core/unittest/test_exec.cpp index 4ab768f255..2ad7f964cd 100644 --- a/internal/core/unittest/test_exec.cpp +++ b/internal/core/unittest/test_exec.cpp @@ -13,14 +13,11 @@ #include #include #include -#include #include #include -#include "query/PlanNode.h" -#include "query/ExecPlanNodeVisitor.h" #include "segcore/SegmentSealed.h" -#include "test_utils/AssertUtils.h" +#include "test_utils/storage_test_utils.h" #include "test_utils/DataGen.h" #include "plan/PlanNode.h" #include "exec/Task.h" @@ -87,25 +84,10 @@ class TaskTest : public testing::TestWithParam { field_map_.insert({"json", json_fid}); schema->set_primary_field_id(str1_fid); - auto segment = CreateSealedSegment(schema); size_t N = 100000; num_rows_ = N; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } + auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data); segment_ = SegmentSealedSPtr(segment.release()); } @@ -181,6 +163,7 @@ TEST_P(TaskTest, CallExprEmpty) { EXPECT_EQ(num_rows, num_rows_); } +// TODO(tiered storage 1): this is slower due to the overhead of RawAt. TEST_P(TaskTest, UnaryExpr) { ::milvus::proto::plan::GenericValue value; value.set_int64_val(-1); diff --git a/internal/core/unittest/test_expr.cpp b/internal/core/unittest/test_expr.cpp index 1b8a8e91a8..502de84e06 100644 --- a/internal/core/unittest/test_expr.cpp +++ b/internal/core/unittest/test_expr.cpp @@ -30,7 +30,6 @@ #include "common/FieldDataInterface.h" #include "common/Json.h" #include "common/JsonCastType.h" -#include "common/LoadInfo.h" #include "common/Types.h" #include "gtest/gtest.h" #include "index/Meta.h" @@ -41,23 +40,19 @@ #include "pb/schema.pb.h" #include "query/Plan.h" #include "query/PlanNode.h" -#include "query/PlanProto.h" #include "query/ExecPlanNodeVisitor.h" #include "segcore/SegmentGrowingImpl.h" #include "simdjson/padded_string.h" -#include "segcore/segment_c.h" #include "storage/FileManager.h" #include "storage/Types.h" #include "storage/Util.h" #include "test_utils/DataGen.h" #include "test_utils/GenExprProto.h" +#include "test_utils/storage_test_utils.h" #include "index/IndexFactory.h" -#include "exec/expression/Expr.h" #include "exec/Task.h" #include "exec/expression/function/FunctionFactory.h" #include "expr/ITypeExpr.h" -#include "index/BitmapIndex.h" -#include "index/InvertedIndexTantivy.h" #include "mmap/Types.h" using namespace milvus; @@ -3478,22 +3473,7 @@ TEST_P(ExprTest, test_term_pk_with_sorted) { schema, nullptr, 1, SegcoreConfig::default_config(), false, true); int N = 100000; auto raw_data = DataGen(schema, N); - - // load field data - auto fields = schema->get_fields(); - - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); std::vector retrieve_ints; for (int i = 0; i < 10; ++i) { @@ -3573,19 +3553,7 @@ TEST_P(ExprTest, TestSealedSegmentGetBatchSize) { auto seg = CreateSealedSegment(schema); size_t N = 1000; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); auto build_expr = [&](enum DataType type) -> expr::TypedExprPtr { @@ -3724,19 +3692,7 @@ TEST_P(ExprTest, TestReorder) { auto seg = CreateSealedSegment(schema); size_t N = 1000; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -3870,19 +3826,7 @@ TEST_P(ExprTest, TestCompareExprNullable) { auto seg = CreateSealedSegment(schema); size_t N = 1000; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); auto build_expr = [&](enum DataType type) -> expr::TypedExprPtr { @@ -4027,19 +3971,7 @@ TEST_P(ExprTest, TestCompareExprNullable2) { auto seg = CreateSealedSegment(schema); size_t N = 1000; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); auto build_expr = [&](enum DataType type) -> expr::TypedExprPtr { @@ -4178,19 +4110,7 @@ TEST_P(ExprTest, TestMutiInConvert) { auto seg = CreateSealedSegment(schema); size_t N = 1000; auto raw_data = DataGen(schema, N); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -4267,19 +4187,7 @@ TEST(Expr, TestExprPerformance) { auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); enum ExprType { UnaryRangeExpr = 0, @@ -4646,19 +4554,7 @@ TEST(Expr, TestExprNOT) { valid_data_double = raw_data.get_col_valid(double_fid); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); enum ExprType { UnaryRangeExpr = 0, @@ -5004,22 +4900,7 @@ TEST_P(ExprTest, test_term_pk) { auto seg = CreateSealedSegment(schema); int N = 1000; auto raw_data = DataGen(schema, N); - - // load field data - auto fields = schema->get_fields(); - - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); std::vector retrieve_ints; for (int i = 0; i < 10; ++i) { @@ -5151,19 +5032,7 @@ TEST_P(ExprTest, TestConjuctExpr) { int N = 1000; auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); auto build_expr = [&](int l, int r) -> expr::TypedExprPtr { @@ -5241,20 +5110,8 @@ TEST_P(ExprTest, TestConjuctExprNullable) { auto seg = CreateSealedSegment(schema); int N = 1000; auto raw_data = DataGen(schema, N); - // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - info.arrow_reader_channel->push( - storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta))); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); auto build_expr = [&](int l, int r) -> expr::TypedExprPtr { @@ -5330,19 +5187,7 @@ TEST_P(ExprTest, TestUnaryBenchTest) { auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -5403,19 +5248,7 @@ TEST_P(ExprTest, TestBinaryRangeBenchTest) { auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -5484,19 +5317,7 @@ TEST_P(ExprTest, TestLogicalUnaryBenchTest) { auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -5560,19 +5381,7 @@ TEST_P(ExprTest, TestBinaryLogicalBenchTest) { auto raw_data = DataGen(schema, N); // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -5645,21 +5454,7 @@ TEST_P(ExprTest, TestBinaryArithOpEvalRangeBenchExpr) { auto seg = CreateSealedSegment(schema); int N = 1000; auto raw_data = DataGen(schema, N); - - // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray( - N, &field_data, fields.at(FieldId(field_id)))); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -5768,20 +5563,7 @@ TEST(Expr, TestExprNull) { FixedVector valid_data_all_true(N, true); - // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray( - N, &field_data, fields.at(FieldId(field_id)))); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); auto build_nullable_expr = [&](DataType data_type, NullExprType op) -> expr::TypedExprPtr { @@ -5958,21 +5740,7 @@ TEST_P(ExprTest, TestCompareExprBenchTest) { int N = 1000; auto raw_data = DataGen(schema, N); - // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray( - N, &field_data, fields.at(FieldId(field_id)))); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); query::ExecPlanNodeVisitor visitor(*seg, MAX_TIMESTAMP); @@ -6030,20 +5798,7 @@ TEST_P(ExprTest, TestRefactorExprs) { int N = 1000; auto raw_data = DataGen(schema, N); - // load field data - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N, "/tmp/a"); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray( - N, &field_data, fields.at(FieldId(field_id)))); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - seg->LoadFieldData(FieldId(field_id), info); - } + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); enum ExprType { UnaryRangeExpr = 0, @@ -6069,13 +5824,6 @@ TEST_P(ExprTest, TestRefactorExprs) { } case TermExprImpl: { std::vector retrieve_ints; - // for (int i = 0; i < n; ++i) { - // retrieve_ints.push_back("xxxxxx" + std::to_string(i % 10)); - // } - // return std::make_shared>( - // ColumnInfo(str1_fid, DataType::VARCHAR), - // retrieve_ints, - // proto::plan::GenericValue::ValCase::kStringVal); for (int i = 0; i < n; ++i) { proto::plan::GenericValue val; val.set_float_val(i); @@ -10847,6 +10595,7 @@ TEST_P(ExprTest, TestBinaryArithOpEvalRangeWithScalarSortIndex) { auto seg = CreateSealedSegment(schema); int N = 1000; auto raw_data = DataGen(schema, N); + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); segcore::LoadIndexInfo load_index_info; // load index for int8 field @@ -11579,6 +11328,7 @@ TEST_P(ExprTest, TestBinaryArithOpEvalRangeWithScalarSortIndexNullable) { auto seg = CreateSealedSegment(schema); int N = 1000; auto raw_data = DataGen(schema, N); + LoadGeneratedDataIntoSegment(raw_data, seg.get(), true); segcore::LoadIndexInfo load_index_info; auto i8_valid_data = raw_data.get_col_valid(i8_nullable_fid); @@ -16634,11 +16384,11 @@ TYPED_TEST(JsonIndexTestFixture, TestJsonIndexUnaryExpr) { load_index_info.index_params = {{JSON_PATH, this->json_path}}; seg->LoadIndex(load_index_info); - auto json_field_data_info = FieldDataInfo( - json_fid.get(), - N, - {storage::ConvertFieldDataToArrowDataWrapper(json_field)}); - seg->LoadFieldData(json_fid, json_field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog( + 1, 1, 1, json_fid.get(), {json_field}, cm); + seg->LoadFieldData(load_info); auto unary_expr = std::make_shared( expr::ColumnInfo(json_fid, DataType::JSON, {this->json_path.substr(1)}), @@ -16768,14 +16518,11 @@ TEST(JsonIndexTest, TestJsonNotEqualExpr) { load_index_info.index_params = {{JSON_PATH, "/a"}}; seg->LoadIndex(load_index_info); - auto json_field_data_info = - FieldDataInfo(json_fid.get(), 2 * json_strs.size()); - json_field_data_info.arrow_reader_channel->push( - storage::ConvertFieldDataToArrowDataWrapper(json_field)); - json_field_data_info.arrow_reader_channel->push( - storage::ConvertFieldDataToArrowDataWrapper(json_field2)); - json_field_data_info.arrow_reader_channel->close(); - seg->LoadFieldData(json_fid, json_field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog( + 1, 1, 1, json_fid.get(), {json_field, json_field2}, cm); + seg->LoadFieldData(load_info); proto::plan::GenericValue val; val.set_int64_val(1); @@ -16874,11 +16621,11 @@ TEST_P(JsonIndexExistsTest, TestExistsExpr) { load_index_info.index_params = {{JSON_PATH, json_index_path}}; seg->LoadIndex(load_index_info); - auto json_field_data_info = FieldDataInfo(json_fid.get(), json_strs.size()); - json_field_data_info.arrow_reader_channel->push( - storage::ConvertFieldDataToArrowDataWrapper(json_field)); - json_field_data_info.arrow_reader_channel->close(); - seg->LoadFieldData(json_fid, json_field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog( + 1, 1, 1, json_fid.get(), {json_field}, cm); + seg->LoadFieldData(load_info); for (auto& [nested_path, exists, expect] : test_cases) { BitsetType expect_res; @@ -17052,12 +16799,11 @@ TEST_P(JsonIndexBinaryExprTest, TestBinaryRangeExpr) { load_index_info.index_params = {{JSON_PATH, "/a"}}; seg->LoadIndex(load_index_info); - auto json_field_data_info = FieldDataInfo( - json_fid.get(), - json_strs.size(), - {storage::ConvertFieldDataToArrowDataWrapper(json_field)}); - - seg->LoadFieldData(json_fid, json_field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog( + 1, 1, 1, json_fid.get(), {json_field}, cm); + seg->LoadFieldData(load_info); for (auto& [lower, upper, lower_inclusive, upper_inclusive, result] : test_cases) { diff --git a/internal/core/unittest/test_expr_materialized_view.cpp b/internal/core/unittest/test_expr_materialized_view.cpp index 6837b34570..0f0b9cb614 100644 --- a/internal/core/unittest/test_expr_materialized_view.cpp +++ b/internal/core/unittest/test_expr_materialized_view.cpp @@ -21,17 +21,16 @@ #include "common/FieldDataInterface.h" #include "common/Schema.h" #include "common/Types.h" -#include "expr/ITypeExpr.h" #include "knowhere/comp/index_param.h" #include "knowhere/comp/materialized_view.h" #include "knowhere/config.h" #include "query/Plan.h" #include "query/PlanImpl.h" #include "query/ExecPlanNodeVisitor.h" -#include "plan/PlanNode.h" #include "segcore/SegmentSealed.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" using DataType = milvus::DataType; using Schema = milvus::Schema; @@ -95,18 +94,17 @@ class ExprMaterializedViewTest : public testing::Test { } else { schema->AddDebugField(field_name, data_type); } - data_field_info[data_type].field_id = - schema->get_field_id(FieldName(field_name)).get(); - std::cout << field_name << " with id " - << data_field_info[data_type].field_id << std::endl; + auto field_id = schema->get_field_id(FieldName(field_name)); + if (data_type == DataType::INT64) { + schema->set_primary_field_id(field_id); + } + data_field_info[data_type].field_id = field_id.get(); } // generate data and prepare for search gen_data = std::make_unique( milvus::segcore::DataGen(schema, N)); - segment = milvus::segcore::CreateSealedSegment(schema); - auto fields = schema->get_fields(); - milvus::segcore::SealedLoadFieldData(*gen_data, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, *gen_data); exec_plan_node_visitor = std::make_unique( *segment, milvus::MAX_TIMESTAMP); @@ -222,13 +220,13 @@ class ExprMaterializedViewTest : public testing::Test { } knowhere::MaterializedViewSearchInfo - TranslateThenExecuteWhenMvInolved(const std::string& predicate_str) { + TranslateThenExecuteWhenMvInvolved(const std::string& predicate_str) { auto plan = CreatePlan(predicate_str, true); return ExecutePlan(plan); } knowhere::MaterializedViewSearchInfo - TranslateThenExecuteWhenMvNotInolved(const std::string& predicate_str) { + TranslateThenExecuteWhenMvNotInvolved(const std::string& predicate_str) { auto plan = CreatePlan(predicate_str, false); return ExecutePlan(plan); } @@ -384,7 +382,7 @@ TEST_F(ExprMaterializedViewTest, TestJsonContainsExpr) { InterpolateSingleExpr( R"( elements: op:Contains elements_same_type:true>)", DataType::INT64); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); TestMvExpectDefault(mv); } @@ -401,7 +399,7 @@ TEST_F(ExprMaterializedViewTest, TestInExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); // fmt::print("Predicate: {}\n", predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); @@ -429,7 +427,7 @@ TEST_F(ExprMaterializedViewTest, TestInDuplicatesExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); // fmt::print("Predicate: {}\n", predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); @@ -460,7 +458,7 @@ TEST_F(ExprMaterializedViewTest, TestUnaryLogicalNotInExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -486,7 +484,7 @@ TEST_F(ExprMaterializedViewTest, TestUnaryRangeEqualExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -512,7 +510,7 @@ TEST_F(ExprMaterializedViewTest, TestUnaryRangeNotEqualExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -542,7 +540,7 @@ TEST_F(ExprMaterializedViewTest, TestUnaryRangeCompareExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -569,7 +567,7 @@ TEST_F(ExprMaterializedViewTest, TestInMultipleExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -600,7 +598,7 @@ TEST_F(ExprMaterializedViewTest, TestUnaryLogicalNotInMultipleExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); @@ -648,7 +646,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualAndEqualExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 2); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -696,7 +694,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualAndInExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 2); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -749,7 +747,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualAndNotInExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 2); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -797,7 +795,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualOrEqualExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 2); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -867,7 +865,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualAndInOrEqualExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 3); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -944,7 +942,7 @@ TEST_F(ExprMaterializedViewTest, TestEqualAndNotEqualOrEqualExpr) { > )"; - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 3); EXPECT_EQ(mv.field_id_to_touched_categories_cnt[GetFieldID(c0_data_type)], @@ -971,7 +969,7 @@ TEST_F(ExprMaterializedViewTest, TestBinaryRangeExpr) { > )"; predicate = InterpolateSingleExpr(predicate, data_type); - auto mv = TranslateThenExecuteWhenMvInolved(predicate); + auto mv = TranslateThenExecuteWhenMvInvolved(predicate); ASSERT_EQ(mv.field_id_to_touched_categories_cnt.size(), 1); auto field_id = GetFieldID(data_type); diff --git a/internal/core/unittest/test_group_by.cpp b/internal/core/unittest/test_group_by.cpp index dfe37fa4f2..f1b1483e6a 100644 --- a/internal/core/unittest/test_group_by.cpp +++ b/internal/core/unittest/test_group_by.cpp @@ -19,6 +19,7 @@ #include "segcore/segment_c.h" #include "test_utils/DataGen.h" #include "test_utils/c_api_test_utils.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::query; @@ -28,35 +29,6 @@ using namespace milvus::tracer; const char* METRICS_TYPE = "metric_type"; -void -prepareSegmentSystemFieldData(const std::unique_ptr& segment, - size_t row_count, - GeneratedData& data_set) { - auto field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(data_set.row_ids_.data(), row_count); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo{ - RowFieldID.get(), - row_count, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(RowFieldID, field_data_info); - - field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(data_set.timestamps_.data(), row_count); - - auto timestamp_arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - field_data_info = - FieldDataInfo{TimestampFieldID.get(), - row_count, - std::vector>{ - timestamp_arrow_data_wrapper}}; - segment->LoadFieldData(TimestampFieldID, field_data_info); -} - int GetSearchResultBound(const SearchResult& search_result) { int i = 0; @@ -98,26 +70,11 @@ TEST(GroupBY, SealedIndex) { auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR); auto bool_fid = schema->AddDebugField("bool", DataType::BOOL); schema->set_primary_field_id(str_fid); - auto segment = CreateSealedSegment(schema); size_t N = 50; //2. load raw data auto raw_data = DataGen(schema, N, 42, 0, 8, 10, false, false); - auto fields = schema->get_fields(); - for (const auto& field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - // Assuming 'channel' is not a member of FieldDataInfo, we need to handle it differently - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentSystemFieldData(segment, N, raw_data); + auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data); //3. load index auto vector_data = raw_data.get_col(vec_fid); @@ -453,25 +410,11 @@ TEST(GroupBY, SealedData) { auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR); auto bool_fid = schema->AddDebugField("bool", DataType::BOOL); schema->set_primary_field_id(str_fid); - auto segment = CreateSealedSegment(schema); size_t N = 100; //2. load raw data auto raw_data = DataGen(schema, N, 42, 0, 20, 10, false, false); - auto fields = schema->get_fields(); - for (auto&& field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray( - N, &field_data, fields.at(FieldId(field_id)))); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentSystemFieldData(segment, N, raw_data); + auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data); int topK = 10; int group_size = 5; @@ -541,14 +484,12 @@ TEST(GroupBY, Reduce) { auto schema = std::make_shared(); auto vec_fid = schema->AddDebugField( "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2); - auto int64_fid = schema->AddDebugField("int64", DataType::INT64, true); + auto int64_fid = schema->AddDebugField("int64", DataType::INT64, false); auto fp16_fid = schema->AddDebugField( "fakevec_fp16", DataType::VECTOR_FLOAT16, dim, knowhere::metric::L2); auto bf16_fid = schema->AddDebugField( "fakevec_bf16", DataType::VECTOR_BFLOAT16, dim, knowhere::metric::L2); schema->set_primary_field_id(int64_fid); - auto segment1 = CreateSealedSegment(schema); - auto segment2 = CreateSealedSegment(schema); //1. load raw data size_t N = 100; @@ -561,32 +502,8 @@ TEST(GroupBY, Reduce) { auto raw_data2 = DataGen(schema, N, seed, ts_offset, repeat_count_2, false, false); - auto fields = schema->get_fields(); - //load segment1 raw data - for (auto field_data : raw_data1.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - segment1->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentSystemFieldData(segment1, N, raw_data1); - - //load segment2 raw data - for (auto&& field_data : raw_data2.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - segment2->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentSystemFieldData(segment2, N, raw_data2); + auto segment1 = CreateSealedWithFieldDataLoaded(schema, raw_data1); + auto segment2 = CreateSealedWithFieldDataLoaded(schema, raw_data2); //3. load index auto vector_data_1 = raw_data1.get_col(vec_fid); diff --git a/internal/core/unittest/test_growing.cpp b/internal/core/unittest/test_growing.cpp index f5c144b79c..452116d114 100644 --- a/internal/core/unittest/test_growing.cpp +++ b/internal/core/unittest/test_growing.cpp @@ -42,7 +42,7 @@ TEST(Growing, DeleteCount) { Timestamp begin_ts = 100; auto tss = GenTss(c, begin_ts); auto del_pks = GenPKs(pks.begin(), pks.end()); - auto status = segment->Delete(offset, c, del_pks.get(), tss.data()); + auto status = segment->Delete(c, del_pks.get(), tss.data()); ASSERT_TRUE(status.ok()); auto cnt = segment->get_deleted_count(); @@ -70,11 +70,9 @@ TEST(Growing, RealCount) { // delete half. auto half = c / 2; - auto del_offset1 = 0; auto del_ids1 = GenPKs(pks.begin(), pks.begin() + half); auto del_tss1 = GenTss(half, c); - auto status = - segment->Delete(del_offset1, half, del_ids1.get(), del_tss1.data()); + auto status = segment->Delete(half, del_ids1.get(), del_tss1.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(c - half, segment->get_real_count()); @@ -82,8 +80,7 @@ TEST(Growing, RealCount) { auto del_offset2 = segment->get_deleted_count(); ASSERT_EQ(del_offset2, half); auto del_tss2 = GenTss(half, c + half); - status = - segment->Delete(del_offset2, half, del_ids1.get(), del_tss2.data()); + status = segment->Delete(half, del_ids1.get(), del_tss2.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(c - half, segment->get_real_count()); @@ -92,7 +89,7 @@ TEST(Growing, RealCount) { ASSERT_EQ(del_offset3, half); auto del_ids3 = GenPKs(pks.begin(), pks.end()); auto del_tss3 = GenTss(c, c + half * 2); - status = segment->Delete(del_offset3, c, del_ids3.get(), del_tss3.data()); + status = segment->Delete(c, del_ids3.get(), del_tss3.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(0, segment->get_real_count()); } diff --git a/internal/core/unittest/test_growing_storage_v2.cpp b/internal/core/unittest/test_growing_storage_v2.cpp index 64effccfd1..a0dc7b79e7 100644 --- a/internal/core/unittest/test_growing_storage_v2.cpp +++ b/internal/core/unittest/test_growing_storage_v2.cpp @@ -156,8 +156,8 @@ TEST_F(TestGrowingStorageV2, LoadFieldData) { auto str_fid = schema->AddDebugField("str", milvus::DataType::VARCHAR, true); schema->set_primary_field_id(pk_fid); - auto segment = milvus::segcore::CreateGrowingSegment( - schema, milvus::segcore::empty_index_meta); + auto segment = + milvus::segcore::CreateGrowingSegment(schema, milvus::empty_index_meta); LoadFieldDataInfo load_info; load_info.field_infos = { {0, diff --git a/internal/core/unittest/test_inverted_index.cpp b/internal/core/unittest/test_inverted_index.cpp index 02af03cbf6..1ea46fee67 100644 --- a/internal/core/unittest/test_inverted_index.cpp +++ b/internal/core/unittest/test_inverted_index.cpp @@ -27,13 +27,13 @@ using namespace milvus; namespace milvus::test { auto -gen_field_meta(int64_t collection_id = 1, - int64_t partition_id = 2, - int64_t segment_id = 3, - int64_t field_id = 101, - DataType data_type = DataType::NONE, - DataType element_type = DataType::NONE, - bool nullable = false) -> storage::FieldDataMeta { +gen_field_data_meta(int64_t collection_id = 1, + int64_t partition_id = 2, + int64_t segment_id = 3, + int64_t field_id = 101, + DataType data_type = DataType::NONE, + DataType element_type = DataType::NONE, + bool nullable = false) -> storage::FieldDataMeta { auto meta = storage::FieldDataMeta{ .collection_id = collection_id, .partition_id = partition_id, @@ -109,13 +109,13 @@ test_run() { int64_t index_version = 4000; int64_t lack_binlog_row = 100; - auto field_meta = test::gen_field_meta(collection_id, - partition_id, - segment_id, - field_id, - dtype, - element_type, - nullable); + auto field_meta = test::gen_field_data_meta(collection_id, + partition_id, + segment_id, + field_id, + dtype, + element_type, + nullable); auto index_meta = test::gen_index_meta( segment_id, field_id, index_build_id, index_version); @@ -518,13 +518,13 @@ test_string() { int64_t index_version = 4001; int64_t lack_binlog_row = 100; - auto field_meta = test::gen_field_meta(collection_id, - partition_id, - segment_id, - field_id, - dtype, - DataType::NONE, - nullable); + auto field_meta = test::gen_field_data_meta(collection_id, + partition_id, + segment_id, + field_id, + dtype, + DataType::NONE, + nullable); auto index_meta = test::gen_index_meta( segment_id, field_id, index_build_id, index_version); diff --git a/internal/core/unittest/test_iterative_filter.cpp b/internal/core/unittest/test_iterative_filter.cpp index 048cc2da07..4e050b6789 100644 --- a/internal/core/unittest/test_iterative_filter.cpp +++ b/internal/core/unittest/test_iterative_filter.cpp @@ -14,10 +14,8 @@ #include "query/Plan.h" #include "segcore/reduce_c.h" -#include "segcore/plan_c.h" -#include "segcore/segment_c.h" #include "test_utils/DataGen.h" -#include "test_utils/c_api_test_utils.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::query; @@ -30,33 +28,6 @@ using namespace milvus::tracer; * so we will not cover all expr type here, just some examples */ -void -prepareSegmentFieldData(const std::unique_ptr& segment, - size_t row_count, - GeneratedData& data_set) { - auto field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(data_set.row_ids_.data(), row_count); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo{ - RowFieldID.get(), - row_count, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(RowFieldID, field_data_info); - - field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(data_set.timestamps_.data(), row_count); - auto ts_arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - field_data_info = FieldDataInfo{ - TimestampFieldID.get(), - row_count, - std::vector>{ts_arrow_data_wrapper}}; - segment->LoadFieldData(TimestampFieldID, field_data_info); -} - void CheckFilterSearchResult(const SearchResult& search_result_by_iterative_filter, const SearchResult& search_result_by_pre_filter, @@ -95,25 +66,11 @@ TEST(IterativeFilter, SealedIndex) { auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR); auto bool_fid = schema->AddDebugField("bool", DataType::BOOL); schema->set_primary_field_id(str_fid); - auto segment = CreateSealedSegment(schema); size_t N = 50; //2. load raw data auto raw_data = DataGen(schema, N, 42, 0, 8, 10, false, false); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentFieldData(segment, N, raw_data); + auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data); //3. load index auto vector_data = raw_data.get_col(vec_fid); @@ -318,25 +275,11 @@ TEST(IterativeFilter, SealedData) { auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR); auto bool_fid = schema->AddDebugField("bool", DataType::BOOL); schema->set_primary_field_id(str_fid); - auto segment = CreateSealedSegment(schema); size_t N = 100; //2. load raw data auto raw_data = DataGen(schema, N, 42, 0, 8, 10, false, false); - auto fields = schema->get_fields(); - for (auto field_data : raw_data.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } - prepareSegmentFieldData(segment, N, raw_data); + auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data); int topK = 10; // int8 binaryRange diff --git a/internal/core/unittest/test_kmeans_clustering.cpp b/internal/core/unittest/test_kmeans_clustering.cpp index 81f25b553b..bde5486511 100644 --- a/internal/core/unittest/test_kmeans_clustering.cpp +++ b/internal/core/unittest/test_kmeans_clustering.cpp @@ -182,7 +182,7 @@ test_run() { int64_t nb = 10000; auto field_meta = - gen_field_meta(collection_id, partition_id, segment_id, field_id); + gen_field_data_meta(collection_id, partition_id, segment_id, field_id); auto index_meta = gen_index_meta(segment_id, field_id, index_build_id, index_version); diff --git a/internal/core/unittest/test_query.cpp b/internal/core/unittest/test_query.cpp index 70fd9e3198..55806ca823 100644 --- a/internal/core/unittest/test_query.cpp +++ b/internal/core/unittest/test_query.cpp @@ -13,11 +13,9 @@ #include "pb/schema.pb.h" #include "query/PlanImpl.h" -#include "query/PlanNode.h" -#include "query/ExecPlanNodeVisitor.h" -#include "segcore/SegmentSealed.h" #include "test_utils/AssertUtils.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" using json = nlohmann::json; using namespace milvus; @@ -594,11 +592,7 @@ TEST(Query, DISABLED_FillSegment) { dataset.raw_); return segment; }()); - segments.emplace_back([&] { - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); - return segment; - }()); + segments.emplace_back(CreateSealedWithFieldDataLoaded(schema, dataset)); // add field { diff --git a/internal/core/unittest/test_random_sample.cpp b/internal/core/unittest/test_random_sample.cpp index 7a8c712305..74dba3e9a8 100644 --- a/internal/core/unittest/test_random_sample.cpp +++ b/internal/core/unittest/test_random_sample.cpp @@ -15,7 +15,7 @@ #include "common/Types.h" #include "test_utils/DataGen.h" #include "test_utils/GenExprProto.h" -#include "plan/PlanNode.h" +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::segcore; @@ -51,8 +51,7 @@ TEST_P(RandomSampleTest, SampleOnly) { const int64_t N = 3000; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto plan = std::make_unique(*schema); plan->plan_node_ = std::make_unique(); @@ -82,7 +81,6 @@ TEST_P(RandomSampleTest, SampleWithUnaryFiler) { const int64_t N = 3000; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); auto size = dataset.raw_->mutable_fields_data()->size(); auto i64_col = dataset.raw_->mutable_fields_data() @@ -100,7 +98,7 @@ TEST_P(RandomSampleTest, SampleWithUnaryFiler) { } } - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); milvus::proto::plan::GenericValue val; val.set_int64_val(1); @@ -143,9 +141,7 @@ TEST(RandomSampleTest, SampleWithEmptyInput) { const int64_t N = 3000; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); milvus::proto::plan::GenericValue val; val.set_int64_val(0); diff --git a/internal/core/unittest/test_regex_query.cpp b/internal/core/unittest/test_regex_query.cpp index 0059e025bd..cb852a0e7a 100644 --- a/internal/core/unittest/test_regex_query.cpp +++ b/internal/core/unittest/test_regex_query.cpp @@ -11,24 +11,19 @@ #include #include -#include #include "pb/plan.pb.h" -#include "segcore/segcore_init_c.h" #include "segcore/SegmentSealed.h" #include "segcore/SegmentGrowing.h" #include "segcore/SegmentGrowingImpl.h" #include "pb/schema.pb.h" #include "test_utils/DataGen.h" -#include "index/IndexFactory.h" -#include "query/Plan.h" -#include "knowhere/comp/brute_force.h" #include "test_utils/GenExprProto.h" #include "query/PlanProto.h" #include "query/ExecPlanNodeVisitor.h" #include "index/InvertedIndexTantivy.h" - +#include "test_utils/storage_test_utils.h" using namespace milvus; using namespace milvus::query; using namespace milvus::segcore; @@ -199,7 +194,6 @@ class SealedSegmentRegexQueryTest : public ::testing::Test { void SetUp() override { schema = GenTestSchema(); - seg = CreateSealedSegment(schema); raw_str = { "b\n", "a\n", @@ -238,7 +232,7 @@ class SealedSegmentRegexQueryTest : public ::testing::Test { json_col->at(i) = raw_json[i]; } - SealedLoadFieldData(raw_data, *seg); + seg = CreateSealedWithFieldDataLoaded(schema, raw_data); } void diff --git a/internal/core/unittest/test_retrieve.cpp b/internal/core/unittest/test_retrieve.cpp index d570bdb391..33ea32e23f 100644 --- a/internal/core/unittest/test_retrieve.cpp +++ b/internal/core/unittest/test_retrieve.cpp @@ -14,8 +14,8 @@ #include "common/Types.h" #include "knowhere/comp/index_param.h" #include "test_utils/DataGen.h" +#include "test_utils/storage_test_utils.h" #include "test_utils/GenExprProto.h" -#include "plan/PlanNode.h" using namespace milvus; using namespace milvus::segcore; @@ -61,8 +61,7 @@ TEST_P(RetrieveTest, AutoID) { auto choose = [=](int i) { return i * 3 % N; }; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto i64_col = dataset.get_col(fid_64); auto plan = std::make_unique(*schema); @@ -119,8 +118,7 @@ TEST_P(RetrieveTest, AutoID2) { auto choose = [=](int i) { return i * 3 % N; }; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto i64_col = dataset.get_col(fid_64); auto plan = std::make_unique(*schema); @@ -180,8 +178,7 @@ TEST_P(RetrieveTest, NotExist) { auto choose2 = [=](int i) { return i * 3 % N + 3 * N; }; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto i64_col = dataset.get_col(fid_64); auto plan = std::make_unique(*schema); @@ -290,8 +287,7 @@ TEST_P(RetrieveTest, Limit) { int64_t N = 101; auto dataset = DataGen(schema, N, 42); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto plan = std::make_unique(*schema); proto::plan::GenericValue unary_val; @@ -339,8 +335,7 @@ TEST_P(RetrieveTest, FillEntry) { int64_t N = 101; auto dataset = DataGen(schema, N, 42); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto plan = std::make_unique(*schema); proto::plan::GenericValue unary_val; unary_val.set_int64_val(0); @@ -383,8 +378,7 @@ TEST_P(RetrieveTest, LargeTimestamp) { auto choose = [=](int i) { return i * choose_sep % N; }; uint64_t ts_offset = 100; auto dataset = DataGen(schema, N, 42, ts_offset + 1); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto i64_col = dataset.get_col(fid_64); auto plan = std::make_unique(*schema); @@ -452,8 +446,7 @@ TEST_P(RetrieveTest, Delete) { auto choose = [=](int i) { return i; }; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto i64_col = dataset.get_col(fid_64); auto ts_col = dataset.get_col(fid_ts); @@ -535,10 +528,7 @@ TEST_P(RetrieveTest, Delete) { auto ids = std::make_unique(); ids->mutable_int_id()->mutable_data()->Add(new_pks.begin(), new_pks.end()); std::vector new_timestamps{10, 10, 10, 10, 10, 10}; - auto reserved_offset = segment->get_deleted_count(); - ASSERT_EQ(reserved_offset, row_count); - segment->Delete(reserved_offset, - new_count, + segment->Delete(new_count, ids.get(), reinterpret_cast(new_timestamps.data())); diff --git a/internal/core/unittest/test_scalar_index_creator.cpp b/internal/core/unittest/test_scalar_index_creator.cpp index 134e59b0f5..050cf4fba3 100644 --- a/internal/core/unittest/test_scalar_index_creator.cpp +++ b/internal/core/unittest/test_scalar_index_creator.cpp @@ -12,20 +12,12 @@ #include #include -#include "indexbuilder/index_c.h" -#include "test_utils/DataGen.h" #include "test_utils/indexbuilder_test_utils.h" #define private public #include "indexbuilder/ScalarIndexCreator.h" -#include "indexbuilder/IndexFactory.h" - -TEST(Dummy, Aha) { - std::cout << "aha" << std::endl; -} constexpr int64_t nb = 100; -namespace indexcgo = milvus::proto::indexcgo; namespace schemapb = milvus::proto::schema; using milvus::indexbuilder::ScalarIndexCreatorPtr; using ScalarTestParams = std::pair; diff --git a/internal/core/unittest/test_sealed.cpp b/internal/core/unittest/test_sealed.cpp index 519cf5df8b..4c9c04fdf6 100644 --- a/internal/core/unittest/test_sealed.cpp +++ b/internal/core/unittest/test_sealed.cpp @@ -13,20 +13,14 @@ #include #include -#include "common/Consts.h" -#include "common/FieldMeta.h" +#include "cachinglayer/Utils.h" #include "common/Types.h" -#include "common/Tracer.h" #include "index/IndexFactory.h" #include "knowhere/version.h" - -#include "storage/MmapManager.h" -#include "storage/MinioChunkManager.h" #include "storage/RemoteChunkManagerSingleton.h" -#include "storage/LocalChunkManagerSingleton.h" #include "storage/Util.h" + #include "test_utils/DataGen.h" -#include "test_utils/indexbuilder_test_utils.h" #include "test_utils/storage_test_utils.h" using namespace milvus; @@ -137,7 +131,7 @@ TEST(Sealed, without_predicate) { load_info.index_params["metric_type"] = "L2"; // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(fake_id); sealed_segment->LoadIndex(load_info); @@ -216,7 +210,7 @@ TEST(Sealed, without_search_ef_less_than_limit) { load_info.index_params["metric_type"] = "L2"; // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(fake_id); sealed_segment->LoadIndex(load_info); @@ -337,7 +331,7 @@ TEST(Sealed, with_predicate) { load_info.index_params["metric_type"] = "L2"; // load index for vec field, load raw data for scalar field - auto sealed_segment = SealedCreator(schema, dataset); + auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); sealed_segment->DropFieldData(fake_id); sealed_segment->LoadIndex(load_info); @@ -431,7 +425,7 @@ TEST(Sealed, with_predicate_filter_all) { load_info.index_params["metric_type"] = "L2"; // load index for vec field, load raw data for scalar field - auto ivf_sealed_segment = SealedCreator(schema, dataset); + auto ivf_sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); ivf_sealed_segment->DropFieldData(fake_id); ivf_sealed_segment->LoadIndex(load_info); @@ -466,7 +460,7 @@ TEST(Sealed, with_predicate_filter_all) { hnsw_load_info.index_params["metric_type"] = "L2"; // load index for vec field, load raw data for scalar field - auto hnsw_sealed_segment = SealedCreator(schema, dataset); + auto hnsw_sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset); hnsw_sealed_segment->DropFieldData(fake_id); hnsw_sealed_segment->LoadIndex(hnsw_load_info); @@ -517,33 +511,6 @@ TEST(Sealed, LoadFieldData) { N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT); // auto segment = CreateSealedSegment(schema); - // std::string dsl = R"({ - // "bool": { - // "must": [ - // { - // "range": { - // "double": { - // "GE": -1, - // "LT": 1 - // } - // } - // }, - // { - // "vector": { - // "fakevec": { - // "metric_type": "L2", - // "params": { - // "nprobe": 10 - // }, - // "query": "$0", - // "topk": 5, - // "round_decimal": 3 - // } - // } - // } - // ] - // } - // })"; const char* raw_plan = R"(vector_anns: < field_id: 100 predicates: < @@ -581,7 +548,7 @@ TEST(Sealed, LoadFieldData) { ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp)); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); segment->Search(plan.get(), ph_group.get(), timestamp); segment->DropFieldData(fakevec_id); @@ -624,50 +591,52 @@ TEST(Sealed, LoadFieldData) { auto valid7 = dataset.get_col_valid(int64_nullable_id); auto valid8 = dataset.get_col_valid(double_nullable_id); auto valid9 = dataset.get_col_valid(str_nullable_id); - ASSERT_EQ(chunk_span1.valid_data(), nullptr); - ASSERT_EQ(chunk_span2.valid_data(), nullptr); - ASSERT_EQ(chunk_span3.second.size(), 0); + ASSERT_EQ(chunk_span1.get().valid_data(), nullptr); + ASSERT_EQ(chunk_span2.get().valid_data(), nullptr); + ASSERT_EQ(chunk_span3.get().second.size(), 0); for (int i = 0; i < N; ++i) { - if (chunk_span1.valid_data() == nullptr || - chunk_span1.valid_data()[i]) { - ASSERT_EQ(chunk_span1.data()[i], ref1[i]); + if (chunk_span1.get().valid_data() == nullptr || + chunk_span1.get().valid_data()[i]) { + ASSERT_EQ(chunk_span1.get().data()[i], ref1[i]); } - if (chunk_span2.valid_data() == nullptr || - chunk_span2.valid_data()[i]) { - ASSERT_EQ(chunk_span2.data()[i], ref2[i]); + if (chunk_span2.get().valid_data() == nullptr || + chunk_span2.get().valid_data()[i]) { + ASSERT_EQ(chunk_span2.get().data()[i], ref2[i]); } - if (chunk_span3.second.size() == 0 || chunk_span3.second[i]) { - ASSERT_EQ(chunk_span3.first[i], ref3[i]); + if (chunk_span3.get().second.size() == 0 || + chunk_span3.get().second[i]) { + ASSERT_EQ(chunk_span3.get().first[i], ref3[i]); } - if (chunk_span4.valid_data() == nullptr || - chunk_span4.valid_data()[i]) { - ASSERT_EQ(chunk_span4.data()[i], ref4[i]); + if (chunk_span4.get().valid_data() == nullptr || + chunk_span4.get().valid_data()[i]) { + ASSERT_EQ(chunk_span4.get().data()[i], ref4[i]); } - if (chunk_span5.valid_data() == nullptr || - chunk_span5.valid_data()[i]) { - ASSERT_EQ(chunk_span5.data()[i], ref5[i]); + if (chunk_span5.get().valid_data() == nullptr || + chunk_span5.get().valid_data()[i]) { + ASSERT_EQ(chunk_span5.get().data()[i], ref5[i]); } - if (chunk_span6.valid_data() == nullptr || - chunk_span6.valid_data()[i]) { - ASSERT_EQ(chunk_span6.data()[i], ref6[i]); + if (chunk_span6.get().valid_data() == nullptr || + chunk_span6.get().valid_data()[i]) { + ASSERT_EQ(chunk_span6.get().data()[i], ref6[i]); } - if (chunk_span7.valid_data() == nullptr || - chunk_span7.valid_data()[i]) { - ASSERT_EQ(chunk_span7.data()[i], ref7[i]); + if (chunk_span7.get().valid_data() == nullptr || + chunk_span7.get().valid_data()[i]) { + ASSERT_EQ(chunk_span7.get().data()[i], ref7[i]); } - if (chunk_span8.valid_data() == nullptr || - chunk_span8.valid_data()[i]) { - ASSERT_EQ(chunk_span8.data()[i], ref8[i]); + if (chunk_span8.get().valid_data() == nullptr || + chunk_span8.get().valid_data()[i]) { + ASSERT_EQ(chunk_span8.get().data()[i], ref8[i]); } - if (chunk_span9.second.size() == 0 || chunk_span9.second[i]) { - ASSERT_EQ(chunk_span9.first[i], ref9[i]); + if (chunk_span9.get().second.size() == 0 || + chunk_span9.get().second[i]) { + ASSERT_EQ(chunk_span9.get().first[i], ref9[i]); } - ASSERT_EQ(chunk_span4.valid_data()[i], valid4[i]); - ASSERT_EQ(chunk_span5.valid_data()[i], valid5[i]); - ASSERT_EQ(chunk_span6.valid_data()[i], valid6[i]); - ASSERT_EQ(chunk_span7.valid_data()[i], valid7[i]); - ASSERT_EQ(chunk_span8.valid_data()[i], valid8[i]); - ASSERT_EQ(chunk_span9.second[i], valid9[i]); + ASSERT_EQ(chunk_span4.get().valid_data()[i], valid4[i]); + ASSERT_EQ(chunk_span5.get().valid_data()[i], valid5[i]); + ASSERT_EQ(chunk_span6.get().valid_data()[i], valid6[i]); + ASSERT_EQ(chunk_span7.get().valid_data()[i], valid7[i]); + ASSERT_EQ(chunk_span8.get().valid_data()[i], valid8[i]); + ASSERT_EQ(chunk_span9.get().second[i], valid9[i]); } auto sr = segment->Search(plan.get(), ph_group.get(), timestamp); @@ -705,33 +674,6 @@ TEST(Sealed, ClearData) { N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT); auto segment = CreateSealedSegment(schema); - // std::string dsl = R"({ - // "bool": { - // "must": [ - // { - // "range": { - // "double": { - // "GE": -1, - // "LT": 1 - // } - // } - // }, - // { - // "vector": { - // "fakevec": { - // "metric_type": "L2", - // "params": { - // "nprobe": 10 - // }, - // "query": "$0", - // "topk": 5, - // "round_decimal": 3 - // } - // } - // } - // ] - // } - // })"; const char* raw_plan = R"(vector_anns: < field_id: 100 predicates: < @@ -769,7 +711,7 @@ TEST(Sealed, ClearData) { ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp)); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); segment->Search(plan.get(), ph_group.get(), timestamp); segment->DropFieldData(fakevec_id); @@ -791,11 +733,11 @@ TEST(Sealed, ClearData) { auto ref1 = dataset.get_col(counter_id); auto ref2 = dataset.get_col(double_id); auto ref3 = dataset.get_col(str_id)->scalars().string_data().data(); - ASSERT_EQ(chunk_span3.second.size(), 0); + ASSERT_EQ(chunk_span3.get().second.size(), 0); for (int i = 0; i < N; ++i) { - ASSERT_EQ(chunk_span1[i], ref1[i]); - ASSERT_EQ(chunk_span2[i], ref2[i]); - ASSERT_EQ(chunk_span3.first[i], ref3[i]); + ASSERT_EQ(chunk_span1.get()[i], ref1[i]); + ASSERT_EQ(chunk_span2.get()[i], ref2[i]); + ASSERT_EQ(chunk_span3.get().first[i], ref3[i]); } auto sr = segment->Search(plan.get(), ph_group.get(), timestamp); @@ -873,7 +815,7 @@ TEST(Sealed, LoadFieldDataMmap) { ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp)); - SealedLoadFieldData(dataset, *segment, {}, true); + segment = CreateSealedWithFieldDataLoaded(schema, dataset, true); segment->Search(plan.get(), ph_group.get(), timestamp); segment->DropFieldData(fakevec_id); @@ -895,11 +837,11 @@ TEST(Sealed, LoadFieldDataMmap) { auto ref1 = dataset.get_col(counter_id); auto ref2 = dataset.get_col(double_id); auto ref3 = dataset.get_col(str_id)->scalars().string_data().data(); - ASSERT_EQ(chunk_span3.second.size(), 0); + ASSERT_EQ(chunk_span3.get().second.size(), 0); for (int i = 0; i < N; ++i) { - ASSERT_EQ(chunk_span1[i], ref1[i]); - ASSERT_EQ(chunk_span2[i], ref2[i]); - ASSERT_EQ(chunk_span3.first[i], ref3[i]); + ASSERT_EQ(chunk_span1.get()[i], ref1[i]); + ASSERT_EQ(chunk_span2.get()[i], ref2[i]); + ASSERT_EQ(chunk_span3.get().first[i], ref3[i]); } auto sr = segment->Search(plan.get(), ph_group.get(), timestamp); @@ -918,20 +860,7 @@ TEST(Sealed, LoadPkScalarIndex) { schema->set_primary_field_id(pk_id); auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - auto fields = schema->get_fields(); - for (auto field_data : dataset.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - - auto info = FieldDataInfo(field_data.field_id(), N); - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(N, &field_data, field_meta)); - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - segment->LoadFieldData(FieldId(field_id), info); - } + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); LoadIndexInfo pk_index; pk_index.field_id = pk_id.get(); @@ -961,34 +890,6 @@ TEST(Sealed, LoadScalarIndex) { auto indexing = GenVecIndexing( N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT); - auto segment = CreateSealedSegment(schema); - // std::string dsl = R"({ - // "bool": { - // "must": [ - // { - // "range": { - // "double": { - // "GE": -1, - // "LT": 1 - // } - // } - // }, - // { - // "vector": { - // "fakevec": { - // "metric_type": "L2", - // "params": { - // "nprobe": 10 - // }, - // "query": "$0", - // "topk": 5, - // "round_decimal": 3 - // } - // } - // } - // ] - // } - // })"; const char* raw_plan = R"(vector_anns: < field_id: 100 predicates: < @@ -1024,36 +925,11 @@ TEST(Sealed, LoadScalarIndex) { auto ph_group = ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); - LoadFieldDataInfo row_id_info; - FieldMeta row_id_field_meta( - FieldName("RowID"), RowFieldID, DataType::INT64, false, std::nullopt); - auto field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(dataset.row_ids_.data(), N); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo{ - RowFieldID.get(), - N, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(RowFieldID, field_data_info); - - LoadFieldDataInfo ts_info; - FieldMeta ts_field_meta(FieldName("Timestamp"), - TimestampFieldID, - DataType::INT64, - false, - std::nullopt); - field_data = - std::make_shared>(DataType::INT64, false); - field_data->FillFieldData(dataset.timestamps_.data(), N); - auto ts_arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto ts_field_data_info = FieldDataInfo{ - TimestampFieldID.get(), - N, - std::vector>{ts_arrow_data_wrapper}}; - segment->LoadFieldData(TimestampFieldID, ts_field_data_info); + auto segment = CreateSealedWithFieldDataLoaded( + schema, + dataset, + false, + GetExcludedFieldIds(schema, {{0, 1, counter_id.get()}})); LoadIndexInfo vec_info; vec_info.field_id = fakevec_id.get(); @@ -1146,7 +1022,7 @@ TEST(Sealed, Delete) { ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp)); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); int64_t row_count = 5; std::vector pks{1, 2, 3, 4, 5}; @@ -1168,10 +1044,7 @@ TEST(Sealed, Delete) { new_ids->mutable_int_id()->mutable_data()->Add(new_pks.begin(), new_pks.end()); std::vector new_timestamps{10, 10, 10}; - auto reserved_offset = segment->get_deleted_count(); - ASSERT_EQ(reserved_offset, row_count); - segment->Delete(reserved_offset, - new_count, + segment->Delete(new_count, new_ids.get(), reinterpret_cast(new_timestamps.data())); } @@ -1231,7 +1104,7 @@ TEST(Sealed, OverlapDelete) { ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp)); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); int64_t row_count = 5; std::vector pks{1, 2, 3, 4, 5}; @@ -1313,22 +1186,24 @@ TEST(Sealed, BF) { size_t N = 100000; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); std::cout << fake_id.get() << std::endl; - SealedLoadFieldData(dataset, *segment, {fake_id.get()}); + auto segment = CreateSealedWithFieldDataLoaded( + schema, dataset, false, {fake_id.get()}); auto vec_data = GenRandomFloatVecs(N, dim); auto field_data = storage::CreateFieldData(DataType::VECTOR_FLOAT, false, dim); field_data->FillFieldData(vec_data.data(), N); - auto fake_vec_arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = - FieldDataInfo{fake_id.get(), - N, - std::vector>{ - fake_vec_arrow_data_wrapper}}; - segment->LoadFieldData(fake_id, field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + fake_id.get(), + {field_data}, + cm); + + segment->LoadFieldData(load_info); auto topK = 1; auto fmt = boost::format(R"(vector_anns: < @@ -1373,21 +1248,25 @@ TEST(Sealed, BF_Overflow) { size_t N = 10; auto dataset = DataGen(schema, N); - auto segment = CreateSealedSegment(schema); - std::cout << fake_id.get() << std::endl; - SealedLoadFieldData(dataset, *segment, {fake_id.get()}); + auto segment = CreateSealedWithFieldDataLoaded( + schema, + dataset, + false, + GetExcludedFieldIds(schema, {0, 1, i64_fid.get()})); auto vec_data = GenMaxFloatVecs(N, dim); auto field_data = storage::CreateFieldData(DataType::VECTOR_FLOAT, false, dim); field_data->FillFieldData(vec_data.data(), N); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo{ - fake_id.get(), - N, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(fake_id, field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto vec_load_info = PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + fake_id.get(), + {field_data}, + cm); + segment->LoadFieldData(vec_load_info); auto topK = 1; auto fmt = boost::format(R"(vector_anns: < @@ -1415,7 +1294,7 @@ TEST(Sealed, BF_Overflow) { auto result = segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP); auto ves = SearchResultToVector(*result); for (int i = 0; i < num_queries; ++i) { - EXPECT_EQ(ves[0].first, -1); + EXPECT_EQ(ves[i].first, -1); } } @@ -1424,41 +1303,44 @@ TEST(Sealed, DeleteCount) { auto schema = std::make_shared(); auto pk = schema->AddDebugField("pk", DataType::INT64); schema->set_primary_field_id(pk); - auto segment = CreateSealedSegment(schema); - segment->get_insert_record().seal_pks(); + // empty segment + size_t N = 10; + + auto dataset = DataGen(schema, N); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); + + auto insert_record_slot = segment->get_insert_record_slot(); + auto ca = milvus::cachinglayer::SemiInlineGet( + insert_record_slot->PinCells({0})); + auto insert_record = ca->get_cell_of(0); + insert_record->seal_pks(); int64_t c = 10; - auto offset = segment->get_deleted_count(); - ASSERT_EQ(offset, 0); + ASSERT_EQ(segment->get_deleted_count(), 0); Timestamp begin_ts = 100; auto tss = GenTss(c, begin_ts); - auto pks = GenPKs(c, 0); - auto status = segment->Delete(offset, c, pks.get(), tss.data()); + auto pks = GenPKs(c, N); + auto status = segment->Delete(c, pks.get(), tss.data()); ASSERT_TRUE(status.ok()); - auto cnt = segment->get_deleted_count(); - ASSERT_EQ(cnt, 0); + ASSERT_EQ(segment->get_deleted_count(), 0); } { auto schema = std::make_shared(); auto pk = schema->AddDebugField("pk", DataType::INT64); schema->set_primary_field_id(pk); - auto segment = CreateSealedSegment(schema); int64_t c = 10; auto dataset = DataGen(schema, c); auto pks = dataset.get_col(pk); - SealedLoadFieldData(dataset, *segment); - - auto offset = segment->get_deleted_count(); - ASSERT_EQ(offset, 0); + auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); auto iter = std::max_element(pks.begin(), pks.end()); auto delete_pks = GenPKs(c, *iter); Timestamp begin_ts = 100; auto tss = GenTss(c, begin_ts); - auto status = segment->Delete(offset, c, delete_pks.get(), tss.data()); + auto status = segment->Delete(c, delete_pks.get(), tss.data()); ASSERT_TRUE(status.ok()); // 9 of element should be filtered. @@ -1478,37 +1360,29 @@ TEST(Sealed, RealCount) { int64_t c = 10; auto dataset = DataGen(schema, c); auto pks = dataset.get_col(pk); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); // no delete. ASSERT_EQ(c, segment->get_real_count()); // delete half. auto half = c / 2; - auto del_offset1 = segment->get_deleted_count(); - ASSERT_EQ(del_offset1, 0); auto del_ids1 = GenPKs(pks.begin(), pks.begin() + half); auto del_tss1 = GenTss(half, c); - auto status = - segment->Delete(del_offset1, half, del_ids1.get(), del_tss1.data()); + auto status = segment->Delete(half, del_ids1.get(), del_tss1.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(c - half, segment->get_real_count()); // delete duplicate. - auto del_offset2 = segment->get_deleted_count(); - ASSERT_EQ(del_offset2, half); auto del_tss2 = GenTss(half, c + half); - status = - segment->Delete(del_offset2, half, del_ids1.get(), del_tss2.data()); + status = segment->Delete(half, del_ids1.get(), del_tss2.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(c - half, segment->get_real_count()); // delete all. - auto del_offset3 = segment->get_deleted_count(); - ASSERT_EQ(del_offset3, half); auto del_ids3 = GenPKs(pks.begin(), pks.end()); auto del_tss3 = GenTss(c, c + half * 2); - status = segment->Delete(del_offset3, c, del_ids3.get(), del_tss3.data()); + status = segment->Delete(c, del_ids3.get(), del_tss3.data()); ASSERT_TRUE(status.ok()); ASSERT_EQ(0, segment->get_real_count()); } @@ -1560,323 +1434,6 @@ TEST(Sealed, GetVector) { } } -TEST(Sealed, GetVectorFromChunkCache) { - auto dim = 16; - auto topK = 5; - auto N = ROW_COUNT; - auto metric_type = knowhere::metric::L2; - auto index_type = knowhere::IndexEnum::INDEX_FAISS_IVFPQ; - - auto file_name = std::string( - "sealed_test_get_vector_from_chunk_cache/insert_log/1/101/1000000"); - - auto sc = milvus::storage::MmapConfig{}; - milvus::storage::MmapManager::GetInstance().Init(sc); - - auto schema = std::make_shared(); - auto fakevec_id = schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, dim, metric_type); - auto counter_id = schema->AddDebugField("counter", DataType::INT64); - auto double_id = schema->AddDebugField("double", DataType::DOUBLE); - auto nothing_id = schema->AddDebugField("nothing", DataType::INT32); - auto str_id = schema->AddDebugField("str", DataType::VARCHAR); - schema->AddDebugField("int8", DataType::INT8); - schema->AddDebugField("int16", DataType::INT16); - schema->AddDebugField("float", DataType::FLOAT); - schema->set_primary_field_id(counter_id); - - auto dataset = DataGen(schema, N); - auto field_data_meta = - milvus::storage::FieldDataMeta{1, 2, 3, fakevec_id.get()}; - auto field_meta = milvus::FieldMeta(milvus::FieldName("facevec"), - fakevec_id, - milvus::DataType::VECTOR_FLOAT, - dim, - metric_type, - false, - std::nullopt); - - auto rcm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() - .GetRemoteChunkManager(); - auto data = dataset.get_col(fakevec_id); - auto data_slices = std::vector{data.data()}; - auto slice_sizes = std::vector{static_cast(N)}; - auto slice_names = std::vector{file_name}; - PutFieldData(rcm.get(), - data_slices, - slice_sizes, - slice_names, - field_data_meta, - field_meta); - - auto conf = generate_build_conf(index_type, metric_type); - auto ds = knowhere::GenDataSet(N, dim, data.data()); - auto indexing = std::make_unique>( - index_type, - metric_type, - knowhere::Version::GetCurrentVersion().VersionNumber()); - indexing->BuildWithDataset(ds, conf); - auto segment_sealed = CreateSealedSegment(schema); - - LoadIndexInfo vec_info; - vec_info.field_id = fakevec_id.get(); - vec_info.index = std::move(indexing); - vec_info.index_params["metric_type"] = knowhere::metric::L2; - segment_sealed->LoadIndex(vec_info); - - auto field_binlog_info = - FieldBinlogInfo{fakevec_id.get(), - N, - std::vector{N}, - false, - std::vector{file_name}}; - segment_sealed->AddFieldDataInfoForSealed( - LoadFieldDataInfo{std::map{ - {fakevec_id.get(), field_binlog_info}}}); - - auto segment = - dynamic_cast(segment_sealed.get()); - auto has = segment->HasRawData(vec_info.field_id); - EXPECT_FALSE(has); - - auto ids_ds = GenRandomIds(N); - auto result = - segment->get_vector(fakevec_id, ids_ds->GetIds(), ids_ds->GetRows()); - - auto vector = result.get()->mutable_vectors()->float_vector().data(); - EXPECT_TRUE(vector.size() == data.size()); - for (size_t i = 0; i < N; ++i) { - auto id = ids_ds->GetIds()[i]; - for (size_t j = 0; j < dim; ++j) { - auto expect = data[id * dim + j]; - auto actual = vector[i * dim + j]; - AssertInfo(expect == actual, - fmt::format("expect {}, actual {}", expect, actual)); - } - } - - rcm->Remove(file_name); - auto exist = rcm->Exist(file_name); - Assert(!exist); -} - -TEST(Sealed, GetSparseVectorFromChunkCache) { - auto dim = 16; - auto topK = 5; - auto N = ROW_COUNT; - auto metric_type = knowhere::metric::IP; - // TODO: remove SegmentSealedImpl::TEST_skip_index_for_retrieve_ after - // we have a type of sparse index that doesn't include raw data. - auto index_type = knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX; - - auto file_name = std::string( - "sealed_test_get_vector_from_chunk_cache/insert_log/1/101/1000000"); - - auto lcm = milvus::storage::LocalChunkManagerSingleton::GetInstance() - .GetChunkManager(); - - auto schema = std::make_shared(); - auto fakevec_id = schema->AddDebugField( - "fakevec", DataType::VECTOR_SPARSE_FLOAT, dim, metric_type); - auto counter_id = schema->AddDebugField("counter", DataType::INT64); - auto double_id = schema->AddDebugField("double", DataType::DOUBLE); - auto nothing_id = schema->AddDebugField("nothing", DataType::INT32); - auto str_id = schema->AddDebugField("str", DataType::VARCHAR); - schema->AddDebugField("int8", DataType::INT8); - schema->AddDebugField("int16", DataType::INT16); - schema->AddDebugField("float", DataType::FLOAT); - schema->set_primary_field_id(counter_id); - - auto dataset = DataGen(schema, N); - auto field_data_meta = - milvus::storage::FieldDataMeta{1, 2, 3, fakevec_id.get()}; - auto field_meta = milvus::FieldMeta(milvus::FieldName("fakevec"), - fakevec_id, - milvus::DataType::VECTOR_SPARSE_FLOAT, - dim, - metric_type, - false, - std::nullopt); - - auto data = dataset.get_col>(fakevec_id); - - // write to multiple files for better coverage - auto data_slices = std::vector(); - auto slice_sizes = std::vector(); - auto slice_names = std::vector(); - - const int64_t slice_size = (N + 9) / 10; - for (int64_t i = 0; i < N; i += slice_size) { - int64_t current_slice_size = std::min(slice_size, N - i); - data_slices.push_back(data.data() + i); - slice_sizes.push_back(current_slice_size); - slice_names.push_back(file_name + "_" + std::to_string(i / slice_size)); - } - PutFieldData(lcm.get(), - data_slices, - slice_sizes, - slice_names, - field_data_meta, - field_meta); - - auto conf = generate_build_conf(index_type, metric_type); - auto ds = knowhere::GenDataSet(N, dim, data.data()); - auto indexing = std::make_unique>( - index_type, - metric_type, - knowhere::Version::GetCurrentVersion().VersionNumber()); - indexing->BuildWithDataset(ds, conf); - auto segment_sealed = CreateSealedSegment( - schema, nullptr, -1, SegcoreConfig::default_config(), true); - - LoadIndexInfo vec_info; - vec_info.field_id = fakevec_id.get(); - vec_info.index = std::move(indexing); - vec_info.index_params["metric_type"] = metric_type; - segment_sealed->LoadIndex(vec_info); - - auto field_binlog_info = - FieldBinlogInfo{fakevec_id.get(), N, slice_sizes, false, slice_names}; - segment_sealed->AddFieldDataInfoForSealed( - LoadFieldDataInfo{std::map{ - {fakevec_id.get(), field_binlog_info}}}); - - auto segment = - dynamic_cast(segment_sealed.get()); - - auto ids_ds = GenRandomIds(N); - auto result = - segment->get_vector(fakevec_id, ids_ds->GetIds(), ids_ds->GetRows()); - - auto vector = - result.get()->mutable_vectors()->sparse_float_vector().contents(); - // number of rows - EXPECT_TRUE(vector.size() == data.size()); - auto sparse_rows = SparseBytesToRows(vector, true); - for (size_t i = 0; i < N; ++i) { - auto expect = data[ids_ds->GetIds()[i]]; - auto& actual = sparse_rows[i]; - AssertInfo( - expect.size() == actual.size(), - fmt::format("expect {}, actual {}", expect.size(), actual.size())); - AssertInfo( - memcmp(expect.data(), actual.data(), expect.data_byte_size()) == 0, - "sparse float vector doesn't match"); - } - - for (const auto& name : slice_names) { - lcm->Remove(name); - auto exist = lcm->Exist(name); - Assert(!exist); - } -} - -TEST(Sealed, WarmupChunkCache) { - auto dim = 16; - auto topK = 5; - auto N = ROW_COUNT; - auto metric_type = knowhere::metric::L2; - auto index_type = knowhere::IndexEnum::INDEX_FAISS_IVFPQ; - - auto mmap_dir = "/tmp/mmap"; - auto file_name = std::string( - "sealed_test_get_vector_from_chunk_cache/insert_log/1/101/1000000"); - - auto sc = milvus::storage::MmapConfig{}; - milvus::storage::MmapManager::GetInstance().Init(sc); - - auto schema = std::make_shared(); - auto fakevec_id = schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, dim, metric_type); - auto counter_id = schema->AddDebugField("counter", DataType::INT64); - auto double_id = schema->AddDebugField("double", DataType::DOUBLE); - auto nothing_id = schema->AddDebugField("nothing", DataType::INT32); - auto str_id = schema->AddDebugField("str", DataType::VARCHAR); - schema->AddDebugField("int8", DataType::INT8); - schema->AddDebugField("int16", DataType::INT16); - schema->AddDebugField("float", DataType::FLOAT); - schema->set_primary_field_id(counter_id); - - auto dataset = DataGen(schema, N); - auto field_data_meta = - milvus::storage::FieldDataMeta{1, 2, 3, fakevec_id.get()}; - auto field_meta = milvus::FieldMeta(milvus::FieldName("facevec"), - fakevec_id, - milvus::DataType::VECTOR_FLOAT, - dim, - metric_type, - false, - std::nullopt); - - auto rcm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() - .GetRemoteChunkManager(); - auto data = dataset.get_col(fakevec_id); - auto data_slices = std::vector{data.data()}; - auto slice_sizes = std::vector{static_cast(N)}; - auto slice_names = std::vector{file_name}; - PutFieldData(rcm.get(), - data_slices, - slice_sizes, - slice_names, - field_data_meta, - field_meta); - - auto conf = generate_build_conf(index_type, metric_type); - auto ds = knowhere::GenDataSet(N, dim, data.data()); - auto indexing = std::make_unique>( - index_type, - metric_type, - knowhere::Version::GetCurrentVersion().VersionNumber()); - indexing->BuildWithDataset(ds, conf); - auto segment_sealed = CreateSealedSegment(schema); - - LoadIndexInfo vec_info; - vec_info.field_id = fakevec_id.get(); - vec_info.index = std::move(indexing); - vec_info.index_params["metric_type"] = knowhere::metric::L2; - segment_sealed->LoadIndex(vec_info); - - auto field_binlog_info = - FieldBinlogInfo{fakevec_id.get(), - N, - std::vector{N}, - false, - std::vector{file_name}}; - segment_sealed->AddFieldDataInfoForSealed( - LoadFieldDataInfo{std::map{ - {fakevec_id.get(), field_binlog_info}}}); - - auto segment = - dynamic_cast(segment_sealed.get()); - auto has = segment->HasRawData(vec_info.field_id); - EXPECT_FALSE(has); - - segment_sealed->WarmupChunkCache(FieldId(vec_info.field_id), true); - - auto ids_ds = GenRandomIds(N); - auto result = - segment->get_vector(fakevec_id, ids_ds->GetIds(), ids_ds->GetRows()); - - auto vector = result.get()->mutable_vectors()->float_vector().data(); - EXPECT_TRUE(vector.size() == data.size()); - for (size_t i = 0; i < N; ++i) { - auto id = ids_ds->GetIds()[i]; - for (size_t j = 0; j < dim; ++j) { - auto expect = data[id * dim + j]; - auto actual = vector[i * dim + j]; - AssertInfo(expect == actual, - fmt::format("expect {}, actual {}", expect, actual)); - } - } - - rcm->Remove(file_name); - std::filesystem::remove_all(mmap_dir); - auto exist = rcm->Exist(file_name); - Assert(!exist); - exist = std::filesystem::exists(mmap_dir); - Assert(!exist); -} - TEST(Sealed, LoadArrayFieldData) { auto dim = 16; auto topK = 5; @@ -1924,7 +1481,7 @@ TEST(Sealed, LoadArrayFieldData) { auto ph_group = ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); - SealedLoadFieldData(dataset, *segment); + segment = CreateSealedWithFieldDataLoaded(schema, dataset); segment->Search(plan.get(), ph_group.get(), 1L << 63); auto ids_ds = GenRandomIds(N); @@ -1981,7 +1538,7 @@ TEST(Sealed, LoadArrayFieldDataWithMMap) { auto ph_group = ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); - SealedLoadFieldData(dataset, *segment, {}, true); + segment = CreateSealedWithFieldDataLoaded(schema, dataset, true); segment->Search(plan.get(), ph_group.get(), 1L << 63); } @@ -2266,13 +1823,15 @@ TEST(Sealed, SkipIndexSkipStringRange) { auto string_field_data = storage::CreateFieldData(DataType::VARCHAR, false, 1, N); string_field_data->FillFieldData(strings.data(), N); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(string_field_data); - auto string_field_data_info = FieldDataInfo{ - string_fid.get(), - N, - std::vector>{arrow_data_wrapper}}; - segment->LoadFieldData(string_fid, string_field_data_info); + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + string_fid.get(), + {string_field_data}, + cm); + segment->LoadFieldData(load_info); auto& skip_index = segment->GetSkipIndex(); ASSERT_TRUE(skip_index.CanSkipUnaryRange( string_fid, 0, OpType::Equal, "w")); @@ -2360,7 +1919,8 @@ TEST(Sealed, QueryAllFields) { int64_t dataset_size = 1000; int64_t dim = 128; auto dataset = DataGen(schema, dataset_size); - SealedLoadFieldData(dataset, *segment); + segment_sealed = CreateSealedWithFieldDataLoaded(schema, dataset); + segment = dynamic_cast(segment_sealed.get()); auto bool_values = dataset.get_col(bool_field); auto int8_values = dataset.get_col(int8_field); @@ -2516,7 +2076,8 @@ TEST(Sealed, QueryAllNullableFields) { int64_t dataset_size = 1000; int64_t dim = 128; auto dataset = DataGen(schema, dataset_size); - SealedLoadFieldData(dataset, *segment); + segment_sealed = CreateSealedWithFieldDataLoaded(schema, dataset); + segment = dynamic_cast(segment_sealed.get()); auto bool_values = dataset.get_col(bool_field); auto int8_values = dataset.get_col(int8_field); @@ -2636,7 +2197,7 @@ TEST(Sealed, SearchSortedPk) { int64_t dataset_size = 1000; auto dataset = DataGen(schema, dataset_size, 42, 0, 10); - SealedLoadFieldData(dataset, *segment); + LoadGeneratedDataIntoSegment(dataset, segment); auto pk_values = dataset.get_col(varchar_pk_field); auto offsets = segment->search_pk(PkType(pk_values[100]), Timestamp(99999)); @@ -2644,6 +2205,6 @@ TEST(Sealed, SearchSortedPk) { EXPECT_EQ(100, offsets[0].get()); auto offsets2 = segment->search_pk(PkType(pk_values[100]), int64_t(105)); - EXPECT_EQ(5, offsets2.size()); + EXPECT_EQ(6, offsets2.size()); EXPECT_EQ(100, offsets2[0].get()); } diff --git a/internal/core/unittest/test_segcore.cpp b/internal/core/unittest/test_segcore.cpp index 59d2d36f6a..375194801e 100644 --- a/internal/core/unittest/test_segcore.cpp +++ b/internal/core/unittest/test_segcore.cpp @@ -12,7 +12,6 @@ #include #include #include -#include #include "segcore/SegmentGrowingImpl.h" #include "test_utils/DataGen.h" @@ -91,86 +90,3 @@ TEST(SegmentCoreTest, SmallIndex) { "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); schema->AddDebugField("age", DataType::INT32); } - -TEST(InsertRecordTest, growing_int64_t) { - using namespace milvus::segcore; - auto schema = std::make_shared(); - schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); - auto i64_fid = schema->AddDebugField("age", DataType::INT64); - schema->set_primary_field_id(i64_fid); - auto record = milvus::segcore::InsertRecord(*schema, int64_t(32)); - const int N = 100000; - - for (int i = 1; i <= N; i++) - record.insert_pk(PkType(int64_t(i)), int64_t(i)); - - for (int i = 1; i <= N; i++) { - std::vector offset = - record.search_pk(PkType(int64_t(i)), int64_t(N + 1)); - ASSERT_EQ(offset[0].get(), int64_t(i)); - } -} - -TEST(InsertRecordTest, growing_string) { - using namespace milvus::segcore; - auto schema = std::make_shared(); - schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); - auto i64_fid = schema->AddDebugField("name", DataType::VARCHAR); - schema->set_primary_field_id(i64_fid); - auto record = milvus::segcore::InsertRecord(*schema, int64_t(32)); - const int N = 100000; - - for (int i = 1; i <= N; i++) - record.insert_pk(PkType(std::to_string(i)), int64_t(i)); - - for (int i = 1; i <= N; i++) { - std::vector offset = - record.search_pk(std::to_string(i), int64_t(N + 1)); - ASSERT_EQ(offset[0].get(), int64_t(i)); - } -} - -TEST(InsertRecordTest, sealed_int64_t) { - using namespace milvus::segcore; - auto schema = std::make_shared(); - schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); - auto i64_fid = schema->AddDebugField("age", DataType::INT64); - schema->set_primary_field_id(i64_fid); - auto record = milvus::segcore::InsertRecord(*schema, int64_t(32)); - const int N = 100000; - - for (int i = N; i >= 1; i--) - record.insert_pk(PkType(int64_t(i)), int64_t(i)); - record.seal_pks(); - - for (int i = 1; i <= N; i++) { - std::vector offset = - record.search_pk(PkType(int64_t(i)), int64_t(N + 1)); - ASSERT_EQ(offset[0].get(), int64_t(i)); - } -} - -TEST(InsertRecordTest, sealed_string) { - using namespace milvus::segcore; - auto schema = std::make_shared(); - schema->AddDebugField( - "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); - auto i64_fid = schema->AddDebugField("name", DataType::VARCHAR); - schema->set_primary_field_id(i64_fid); - auto record = milvus::segcore::InsertRecord(*schema, int64_t(32)); - const int N = 100000; - - for (int i = 1; i <= N; i++) - record.insert_pk(PkType(std::to_string(i)), int64_t(i)); - - record.seal_pks(); - - for (int i = 1; i <= N; i++) { - std::vector offset = - record.search_pk(std::to_string(i), int64_t(N + 1)); - ASSERT_EQ(offset[0].get(), int64_t(i)); - } -} \ No newline at end of file diff --git a/internal/core/unittest/test_span.cpp b/internal/core/unittest/test_span.cpp index 8fca5a9c05..6450769bf7 100644 --- a/internal/core/unittest/test_span.cpp +++ b/internal/core/unittest/test_span.cpp @@ -61,21 +61,22 @@ TEST(Span, Naive) { auto begin = chunk_id * size_per_chunk; auto end = std::min((chunk_id + 1) * size_per_chunk, N); auto size_of_chunk = end - begin; - ASSERT_EQ(age_span.valid_data(), nullptr); + ASSERT_EQ(age_span.get().valid_data(), nullptr); for (int i = 0; i < size_of_chunk * 512 / 8; ++i) { - ASSERT_EQ(vec_span.data()[i], vec_ptr[i + begin * 512 / 8]); + ASSERT_EQ(vec_span.get().data()[i], vec_ptr[i + begin * 512 / 8]); } for (int i = 0; i < size_of_chunk; ++i) { - ASSERT_EQ(age_span.data()[i], age_ptr[i + begin]); + ASSERT_EQ(age_span.get().data()[i], age_ptr[i + begin]); } for (int i = 0; i < size_of_chunk; ++i) { - ASSERT_EQ(float_span.data()[i], float_ptr[i + begin * 32]); + ASSERT_EQ(float_span.get().data()[i], float_ptr[i + begin * 32]); } for (int i = 0; i < size_of_chunk; ++i) { - ASSERT_EQ(null_field_span.data()[i], nullable_data_ptr[i + begin]); + ASSERT_EQ(null_field_span.get().data()[i], + nullable_data_ptr[i + begin]); } for (int i = 0; i < size_of_chunk; ++i) { - ASSERT_EQ(null_field_span.valid_data()[i], + ASSERT_EQ(null_field_span.get().valid_data()[i], nullable_valid_data_ptr[i + begin]); } } diff --git a/internal/core/unittest/test_storage.cpp b/internal/core/unittest/test_storage.cpp index 40f80528e8..d030ee7091 100644 --- a/internal/core/unittest/test_storage.cpp +++ b/internal/core/unittest/test_storage.cpp @@ -11,18 +11,16 @@ #include -#include +#include #include #include #include #include "common/EasyAssert.h" #include "storage/LocalChunkManagerSingleton.h" #include "storage/RemoteChunkManagerSingleton.h" +#include "storage/Util.h" #include "storage/storage_c.h" -#define private public -#include "storage/ChunkCache.h" - using namespace std; using namespace milvus; using namespace milvus::storage; @@ -104,9 +102,93 @@ TEST_F(StorageTest, InitRemoteChunkManagerSingleton) { EXPECT_EQ(rcm->GetRootPath(), "/tmp/milvus/remote_data"); } -TEST_F(StorageTest, InitChunkCacheSingleton) { -} - TEST_F(StorageTest, CleanRemoteChunkManagerSingleton) { CleanRemoteChunkManagerSingleton(); -} \ No newline at end of file +} + +class StorageUtilTest : public testing::Test { + public: + StorageUtilTest() = default; + ~StorageUtilTest() { + } + void + SetUp() override { + } +}; + +TEST_F(StorageUtilTest, CreateArrowScalarFromDefaultValue) { + { + FieldMeta field_without_defval( + FieldName("f"), FieldId(100), DataType::INT64, false, std::nullopt); + ASSERT_ANY_THROW( + CreateArrowScalarFromDefaultValue(field_without_defval)); + } + { + DefaultValueType default_value; + default_value.set_int_data(10); + FieldMeta int_field(FieldName("f"), + FieldId(100), + DataType::INT32, + false, + default_value); + auto scalar = CreateArrowScalarFromDefaultValue(int_field); + ASSERT_TRUE(scalar->Equals(*arrow::MakeScalar(int32_t(10)))); + } + { + DefaultValueType default_value; + default_value.set_long_data(10); + FieldMeta long_field(FieldName("f"), + FieldId(100), + DataType::INT64, + false, + default_value); + auto scalar = CreateArrowScalarFromDefaultValue(long_field); + ASSERT_TRUE(scalar->Equals(*arrow::MakeScalar(int64_t(10)))); + } + { + DefaultValueType default_value; + default_value.set_float_data(1.0f); + FieldMeta float_field(FieldName("f"), + FieldId(100), + DataType::FLOAT, + false, + default_value); + auto scalar = CreateArrowScalarFromDefaultValue(float_field); + ASSERT_TRUE(scalar->ApproxEquals(*arrow::MakeScalar(1.0f))); + } + { + DefaultValueType default_value; + default_value.set_double_data(1.0f); + FieldMeta double_field(FieldName("f"), + FieldId(100), + DataType::DOUBLE, + false, + default_value); + auto scalar = CreateArrowScalarFromDefaultValue(double_field); + ASSERT_TRUE(scalar->ApproxEquals(arrow::DoubleScalar(1.0f))); + } + { + DefaultValueType default_value; + default_value.set_bool_data(true); + FieldMeta bool_field( + FieldName("f"), FieldId(100), DataType::BOOL, false, default_value); + auto scalar = CreateArrowScalarFromDefaultValue(bool_field); + ASSERT_TRUE(scalar->Equals(*arrow::MakeScalar(true))); + } + { + DefaultValueType default_value; + default_value.set_string_data("bar"); + FieldMeta varchar_field(FieldName("f"), + FieldId(100), + DataType::VARCHAR, + false, + default_value); + auto scalar = CreateArrowScalarFromDefaultValue(varchar_field); + ASSERT_TRUE(scalar->Equals(*arrow::MakeScalar("bar"))); + } + { + FieldMeta unsupport_field( + FieldName("f"), FieldId(100), DataType::JSON, false, std::nullopt); + ASSERT_ANY_THROW(CreateArrowScalarFromDefaultValue(unsupport_field)); + } +} diff --git a/internal/core/unittest/test_text_match.cpp b/internal/core/unittest/test_text_match.cpp index 8874baa0e7..e0b94d158e 100644 --- a/internal/core/unittest/test_text_match.cpp +++ b/internal/core/unittest/test_text_match.cpp @@ -409,7 +409,6 @@ TEST(TextMatch, GrowingNaiveNullable) { TEST(TextMatch, SealedNaive) { auto schema = GenTestSchema(); - auto seg = CreateSealedSegment(schema, empty_index_meta); std::vector raw_str = {"football, basketball, pingpang", "swimming, football"}; @@ -425,7 +424,7 @@ TEST(TextMatch, SealedNaive) { str_col->at(i) = raw_str[i]; } - SealedLoadFieldData(raw_data, *seg); + auto seg = CreateSealedWithFieldDataLoaded(schema, raw_data); seg->CreateTextIndex(FieldId(101)); { @@ -492,7 +491,6 @@ TEST(TextMatch, SealedNaive) { TEST(TextMatch, SealedNaiveNullable) { auto schema = GenTestSchema({}, true); - auto seg = CreateSealedSegment(schema, empty_index_meta); std::vector raw_str = { "football, basketball, pingpang", "swimming, football", ""}; std::vector raw_str_valid = {true, true, false}; @@ -514,7 +512,7 @@ TEST(TextMatch, SealedNaiveNullable) { str_col_valid->at(i) = raw_str_valid[i]; } - SealedLoadFieldData(raw_data, *seg); + auto seg = CreateSealedWithFieldDataLoaded(schema, raw_data); seg->CreateTextIndex(FieldId(101)); { BitsetType final; @@ -782,7 +780,6 @@ TEST(TextMatch, SealedJieBa) { {"enable_analyzer", "true"}, {"analyzer_params", R"({"tokenizer": "jieba"})"}, }); - auto seg = CreateSealedSegment(schema, empty_index_meta); std::vector raw_str = {"青铜时代", "黄金时代"}; int64_t N = 2; @@ -797,7 +794,7 @@ TEST(TextMatch, SealedJieBa) { str_col->at(i) = raw_str[i]; } - SealedLoadFieldData(raw_data, *seg); + auto seg = CreateSealedWithFieldDataLoaded(schema, raw_data); seg->CreateTextIndex(FieldId(101)); { @@ -866,7 +863,6 @@ TEST(TextMatch, SealedJieBaNullable) { {"analyzer_params", R"({"tokenizer": "jieba"})"}, }, true); - auto seg = CreateSealedSegment(schema, empty_index_meta); std::vector raw_str = {"青铜时代", "黄金时代", ""}; std::vector raw_str_valid = {true, true, false}; @@ -887,7 +883,7 @@ TEST(TextMatch, SealedJieBaNullable) { str_col_valid->at(i) = raw_str_valid[i]; } - SealedLoadFieldData(raw_data, *seg); + auto seg = CreateSealedWithFieldDataLoaded(schema, raw_data); seg->CreateTextIndex(FieldId(101)); { @@ -955,8 +951,9 @@ TEST(TextMatch, SealedJieBaNullable) { } } +// TODO(tiered storage 1): this also fails on master branch. // Test that growing segment loading flushed binlogs will build text match index. -TEST(TextMatch, GrowingLoadData) { +TEST(TextMatch, DISABLED_GrowingLoadData) { int64_t N = 7; auto schema = GenTestSchema({}, true); schema->AddField( @@ -992,13 +989,7 @@ TEST(TextMatch, GrowingLoadData) { auto storage_config = get_default_local_storage_config(); auto cm = storage::CreateChunkManager(storage_config); - auto load_info = PrepareInsertBinlog( - 1, - 2, - 3, - storage_config.root_path + "/" + "test_growing_segment_load_data", - raw_data, - cm); + auto load_info = PrepareInsertBinlog(1, 2, 3, raw_data, cm); auto segment = CreateGrowingSegment(schema, empty_index_meta); auto status = LoadFieldData(segment.get(), &load_info); diff --git a/internal/core/unittest/test_utils.cpp b/internal/core/unittest/test_utils.cpp index 1d4fbcc2ec..b0e496cb97 100644 --- a/internal/core/unittest/test_utils.cpp +++ b/internal/core/unittest/test_utils.cpp @@ -22,6 +22,7 @@ #include "common/EasyAssert.h" #include "common/Types.h" #include "common/Utils.h" +#include "segcore/Record.h" #include "common/Exception.h" #include "knowhere/sparse_utils.h" #include "pb/schema.pb.h" @@ -63,7 +64,7 @@ TEST(Util, GetDeleteBitmap) { schema->set_primary_field_id(i64_fid); auto N = 10; uint64_t seg_id = 101; - InsertRecord insert_record(*schema, N); + InsertRecord insert_record(*schema, N); DeletedRecord delete_record( &insert_record, [&insert_record](const PkType& pk, Timestamp timestamp) { diff --git a/internal/core/unittest/test_utils/AssertUtils.h b/internal/core/unittest/test_utils/AssertUtils.h index 837e44fe76..68b5fd6bcc 100644 --- a/internal/core/unittest/test_utils/AssertUtils.h +++ b/internal/core/unittest/test_utils/AssertUtils.h @@ -15,10 +15,11 @@ #include #include +#include "common/QueryResult.h" #include "common/Types.h" +#include "index/ScalarIndex.h" using milvus::index::ScalarIndex; - namespace { bool diff --git a/internal/core/unittest/test_utils/Constants.h b/internal/core/unittest/test_utils/Constants.h index 3e8858da7d..f853642ccd 100644 --- a/internal/core/unittest/test_utils/Constants.h +++ b/internal/core/unittest/test_utils/Constants.h @@ -10,6 +10,9 @@ // or implied. See the License for the specific language governing permissions and limitations under the License #pragma once + +#include + constexpr int64_t TestChunkSize = 32 * 1024; constexpr char TestLocalPath[] = "/tmp/milvus/local_data/"; constexpr char TestRemotePath[] = "/tmp/milvus/remote_data"; diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index 4417592e29..5e64d6058f 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -36,7 +36,6 @@ #include "PbHelper.h" #include "segcore/collection_c.h" -#include "segcore/SegmentSealed.h" #include "common/Types.h" #include "storage/BinlogReader.h" #include "storage/Event.h" @@ -47,6 +46,10 @@ using boost::algorithm::starts_with; +constexpr int64_t kCollectionID = 1; +constexpr int64_t kPartitionID = 1; +constexpr int64_t kSegmentID = 1; + namespace milvus::segcore { struct GeneratedData { @@ -950,7 +953,7 @@ CreateFieldDataFromDataArray(ssize_t raw_count, int64_t dim) { field_data = storage::CreateFieldData(data_type, true, dim); int byteSize = (raw_count + 7) / 8; - uint8_t* valid_data = new uint8_t[byteSize]; + std::vector valid_data(byteSize); for (int i = 0; i < raw_count; i++) { bool value = raw_valid_data[i]; int byteIndex = i / 8; @@ -961,8 +964,7 @@ CreateFieldDataFromDataArray(ssize_t raw_count, valid_data[byteIndex] &= ~(1 << bitIndex); } } - field_data->FillFieldData(raw_data, valid_data, raw_count); - delete[] valid_data; + field_data->FillFieldData(raw_data, valid_data.data(), raw_count); }; if (field_meta.is_vector()) { @@ -1146,80 +1148,6 @@ CreateFieldDataFromDataArray(ssize_t raw_count, return field_data; } -inline void -SealedLoadFieldData(const GeneratedData& dataset, - SegmentSealed& seg, - const std::set& exclude_fields = {}, - bool with_mmap = false) { - auto row_count = dataset.row_ids_.size(); - { - auto field_data = std::make_shared>( - DataType::INT64, false); - field_data->FillFieldData(dataset.row_ids_.data(), row_count); - - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - - auto field_data_info = FieldDataInfo( - RowFieldID.get(), - row_count, - std::vector>{ - arrow_data_wrapper}); - seg.LoadFieldData(RowFieldID, field_data_info); - } - { - auto field_data = std::make_shared>( - DataType::INT64, false); - field_data->FillFieldData(dataset.timestamps_.data(), row_count); - auto arrow_data_wrapper = - storage::ConvertFieldDataToArrowDataWrapper(field_data); - auto field_data_info = FieldDataInfo( - TimestampFieldID.get(), - row_count, - std::vector>{ - arrow_data_wrapper}); - seg.LoadFieldData(TimestampFieldID, field_data_info); - } - for (auto& iter : dataset.schema_->get_fields()) { - int64_t field_id = iter.first.get(); - if (exclude_fields.find(field_id) != exclude_fields.end()) { - continue; - } - } - auto fields = dataset.schema_->get_fields(); - for (auto& field_data : dataset.raw_->fields_data()) { - int64_t field_id = field_data.field_id(); - if (exclude_fields.find(field_id) != exclude_fields.end()) { - continue; - } - FieldDataInfo info; - if (with_mmap) { - info.mmap_dir_path = "./data/mmap-test"; - } - info.field_id = field_data.field_id(); - info.row_count = row_count; - auto field_meta = fields.at(FieldId(field_id)); - auto arrow_data_wrapper = storage::ConvertFieldDataToArrowDataWrapper( - CreateFieldDataFromDataArray(row_count, &field_data, field_meta)); - - info.arrow_reader_channel->push(arrow_data_wrapper); - info.arrow_reader_channel->close(); - - if (with_mmap) { - seg.MapFieldData(FieldId(field_id), info); - } else { - seg.LoadFieldData(FieldId(field_id), info); - } - } -} - -inline std::unique_ptr -SealedCreator(SchemaPtr schema, const GeneratedData& dataset) { - auto segment = CreateSealedSegment(schema); - SealedLoadFieldData(dataset, *segment); - return segment; -} - inline std::unique_ptr GenVecIndexing(int64_t N, int64_t dim, diff --git a/internal/core/unittest/test_utils/storage_test_utils.h b/internal/core/unittest/test_utils/storage_test_utils.h index bd67a37f37..7dabbde546 100644 --- a/internal/core/unittest/test_utils/storage_test_utils.h +++ b/internal/core/unittest/test_utils/storage_test_utils.h @@ -18,11 +18,16 @@ #include "Constants.h" #include "DataGen.h" +#include "common/Consts.h" +#include "common/IndexMeta.h" #include "common/Types.h" #include "common/LoadInfo.h" +#include "common/Schema.h" +#include "segcore/segment_c.h" #include "storage/Types.h" #include "storage/InsertData.h" -#include "storage/ThreadPools.h" +#include "storage/RemoteChunkManagerSingleton.h" +#include "segcore/SegmentSealed.h" #include using milvus::DataType; @@ -59,15 +64,21 @@ get_default_mmap_config() { return mmap_config; } +// This function uploads generated data into cm, and returns a load info that +// can be used by a segment to load the data. inline LoadFieldDataInfo PrepareInsertBinlog(int64_t collection_id, int64_t partition_id, int64_t segment_id, - const std::string& prefix, const GeneratedData& dataset, - const ChunkManagerPtr cm) { + const ChunkManagerPtr cm, + std::string mmap_dir_path = "", + std::vector excluded_field_ids = {}) { + bool enable_mmap = !mmap_dir_path.empty(); LoadFieldDataInfo load_info; + load_info.mmap_dir_path = mmap_dir_path; auto row_count = dataset.row_ids_.size(); + const std::string prefix = TestRemotePath; auto SaveFieldData = [&](const FieldDataPtr field_data, const std::string& file, @@ -87,78 +98,142 @@ PrepareInsertBinlog(int64_t collection_id, FieldBinlogInfo{field_id, static_cast(row_count), std::vector{int64_t(row_count)}, - false, + enable_mmap, std::vector{file}}); }; - { + auto field_excluded = [&](int64_t field_id) { + return std::find(excluded_field_ids.begin(), + excluded_field_ids.end(), + field_id) != excluded_field_ids.end(); + }; + + if (!field_excluded(RowFieldID.get())) { auto field_data = std::make_shared>( DataType::INT64, false); field_data->FillFieldData(dataset.row_ids_.data(), row_count); - auto path = prefix + "/" + std::to_string(RowFieldID.get()); + auto path = prefix + std::to_string(RowFieldID.get()); SaveFieldData(field_data, path, RowFieldID.get()); } - { + if (!field_excluded(TimestampFieldID.get())) { auto field_data = std::make_shared>( DataType::INT64, false); field_data->FillFieldData(dataset.timestamps_.data(), row_count); - auto path = prefix + "/" + std::to_string(TimestampFieldID.get()); + auto path = prefix + std::to_string(TimestampFieldID.get()); SaveFieldData(field_data, path, TimestampFieldID.get()); } auto fields = dataset.schema_->get_fields(); for (auto& data : dataset.raw_->fields_data()) { int64_t field_id = data.field_id(); - auto field_meta = fields.at(FieldId(field_id)); - auto field_data = milvus::segcore::CreateFieldDataFromDataArray( - row_count, &data, field_meta); - auto path = prefix + "/" + std::to_string(field_id); - SaveFieldData(field_data, path, field_id); + if (!field_excluded(field_id)) { + auto field_meta = fields.at(FieldId(field_id)); + auto field_data = milvus::segcore::CreateFieldDataFromDataArray( + row_count, &data, field_meta); + auto path = prefix + std::to_string(field_id); + SaveFieldData(field_data, path, field_id); + } } return load_info; } -std::map -PutFieldData(milvus::storage::ChunkManager* remote_chunk_manager, - const std::vector& buffers, - const std::vector& element_counts, - const std::vector& object_keys, - FieldDataMeta& field_data_meta, - milvus::FieldMeta& field_meta) { - auto& pool = - milvus::ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE); - std::vector>> futures; - AssertInfo(buffers.size() == element_counts.size(), - "inconsistent size of data slices with slice sizes!"); - AssertInfo(buffers.size() == object_keys.size(), - "inconsistent size of data slices with slice names!"); - - for (int64_t i = 0; i < buffers.size(); ++i) { - futures.push_back( - pool.Submit(milvus::storage::EncodeAndUploadFieldSlice, - remote_chunk_manager, - buffers[i], - element_counts[i], - field_data_meta, - field_meta, - object_keys[i])); +inline LoadFieldDataInfo +PrepareSingleFieldInsertBinlog(int64_t collection_id, + int64_t partition_id, + int64_t segment_id, + int64_t field_id, + std::vector field_datas, + const ChunkManagerPtr cm, + const std::string& mmap_dir_path = "") { + bool enable_mmap = !mmap_dir_path.empty(); + LoadFieldDataInfo load_info; + load_info.mmap_dir_path = mmap_dir_path; + std::vector files; + files.reserve(field_datas.size()); + std::vector row_counts; + row_counts.reserve(field_datas.size()); + int64_t row_count = 0; + for (auto i = 0; i < field_datas.size(); ++i) { + auto& field_data = field_datas[i]; + row_count += field_data->Length(); + auto file = + "./data/test" + std::to_string(field_id) + "/" + std::to_string(i); + files.push_back(file); + row_counts.push_back(field_data->Length()); + auto payload_reader = + std::make_shared(field_data); + auto insert_data = std::make_shared(payload_reader); + FieldDataMeta field_data_meta{ + collection_id, partition_id, segment_id, field_id}; + insert_data->SetFieldDataMeta(field_data_meta); + auto serialized_insert_data = insert_data->serialize_to_remote_file(); + auto serialized_insert_size = serialized_insert_data.size(); + cm->Write(file, serialized_insert_data.data(), serialized_insert_size); } - std::map remote_paths_to_size; - for (auto& future : futures) { - auto res = future.get(); - remote_paths_to_size[res.first] = res.second; - } + load_info.field_infos.emplace( + field_id, + FieldBinlogInfo{field_id, + static_cast(row_count), + row_counts, + enable_mmap, + files}); - milvus::storage::ReleaseArrowUnused(); - return remote_paths_to_size; + return load_info; } +inline void +LoadGeneratedDataIntoSegment(const GeneratedData& dataset, + milvus::segcore::SegmentSealed* segment, + bool with_mmap = false, + std::vector excluded_field_ids = {}) { + std::string mmap_dir_path = with_mmap ? "./data/mmap-test" : ""; + auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto load_info = PrepareInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + dataset, + cm, + mmap_dir_path, + excluded_field_ids); + auto status = LoadFieldData(segment, &load_info); + AssertInfo(status.error_code == milvus::Success, + "Failed to load field data, error: {}", + status.error_msg); +} + +inline std::unique_ptr +CreateSealedWithFieldDataLoaded(milvus::SchemaPtr schema, + const GeneratedData& dataset, + bool with_mmap = false, + std::vector excluded_field_ids = {}) { + auto segment = + milvus::segcore::CreateSealedSegment(schema, milvus::empty_index_meta); + LoadGeneratedDataIntoSegment( + dataset, segment.get(), with_mmap, excluded_field_ids); + return segment; +} + +inline std::vector +GetExcludedFieldIds(milvus::SchemaPtr schema, + std::vector field_ids_to_load) { + auto fields = schema->get_fields(); + std::vector result; + for (auto& field : fields) { + if (std::find(field_ids_to_load.begin(), + field_ids_to_load.end(), + field.first.get()) == field_ids_to_load.end()) { + result.push_back(field.first.get()); + } + } + return result; +} auto -gen_field_meta(int64_t collection_id = 1, - int64_t partition_id = 2, - int64_t segment_id = 3, - int64_t field_id = 101) -> milvus::storage::FieldDataMeta { +gen_field_data_meta(int64_t collection_id = 1, + int64_t partition_id = 2, + int64_t segment_id = 3, + int64_t field_id = 101) -> milvus::storage::FieldDataMeta { return milvus::storage::FieldDataMeta{ .collection_id = collection_id, .partition_id = partition_id, diff --git a/internal/querynodev2/segments/load_field_data_info.go b/internal/querynodev2/segments/load_field_data_info.go index 29eef77a40..b92a42d18b 100644 --- a/internal/querynodev2/segments/load_field_data_info.go +++ b/internal/querynodev2/segments/load_field_data_info.go @@ -101,16 +101,6 @@ func (ld *LoadFieldDataInfo) appendMMapDirPath(dir string) { }).Await() } -func (ld *LoadFieldDataInfo) appendURI(uri string) { - GetDynamicPool().Submit(func() (any, error) { - cURI := C.CString(uri) - defer C.free(unsafe.Pointer(cURI)) - C.SetUri(ld.cLoadFieldDataInfo, cURI) - - return nil, nil - }).Await() -} - func (ld *LoadFieldDataInfo) appendStorageVersion(version int64) { GetDynamicPool().Submit(func() (any, error) { cVersion := C.int64_t(version) diff --git a/internal/querynodev2/segments/segment.go b/internal/querynodev2/segments/segment.go index 90fe5680a4..22c74aeeab 100644 --- a/internal/querynodev2/segments/segment.go +++ b/internal/querynodev2/segments/segment.go @@ -646,19 +646,6 @@ func (s *LocalSegment) RetrieveByOffsets(ctx context.Context, plan *segcore.Retr return retrieveResult, nil } -func (s *LocalSegment) GetFieldDataPath(index *IndexedFieldInfo, offset int64) (dataPath string, offsetInBinlog int64) { - offsetInBinlog = offset - for _, binlog := range index.FieldBinlog.Binlogs { - if offsetInBinlog < binlog.EntriesNum { - dataPath = binlog.GetLogPath() - break - } else { - offsetInBinlog -= binlog.EntriesNum - } - } - return dataPath, offsetInBinlog -} - func (s *LocalSegment) Insert(ctx context.Context, rowIDs []int64, timestamps []typeutil.Timestamp, record *segcorepb.InsertRecord) error { if s.Type() != SegmentTypeGrowing { return fmt.Errorf("unexpected segmentType when segmentInsert, segmentType = %s", s.segmentType.String()) diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go index ae719aa002..16d35b1040 100644 --- a/internal/querynodev2/segments/segment_loader.go +++ b/internal/querynodev2/segments/segment_loader.go @@ -812,7 +812,7 @@ func (loader *segmentLoader) loadSealedSegment(ctx context.Context, loadInfo *qu if err != nil { return err } - if (!typeutil.IsVectorType(field.GetDataType()) && !segment.HasRawData(fieldID)) || field.GetIsPrimaryKey() { + if !segment.HasRawData(fieldID) || field.GetIsPrimaryKey() { log.Info("field index doesn't include raw data, load binlog...", zap.Int64("fieldID", fieldID), zap.String("index", info.IndexInfo.GetIndexName()), diff --git a/internal/querynodev2/segments/segment_loader_test.go b/internal/querynodev2/segments/segment_loader_test.go index 4808a1aae1..44cfe8aa19 100644 --- a/internal/querynodev2/segments/segment_loader_test.go +++ b/internal/querynodev2/segments/segment_loader_test.go @@ -83,6 +83,8 @@ func (suite *SegmentLoaderSuite) SetupTest() { suite.manager = NewManager() suite.loader = NewLoader(ctx, suite.manager, suite.chunkManager) initcore.InitRemoteChunkManager(paramtable.Get()) + initcore.InitLocalChunkManager(suite.rootPath) + initcore.InitMmapManager(paramtable.Get()) // Data suite.schema = mock_segcore.GenTestCollectionSchema("test", schemapb.DataType_Int64, false) diff --git a/internal/querynodev2/server.go b/internal/querynodev2/server.go index cdadeebbf6..455a03f502 100644 --- a/internal/querynodev2/server.go +++ b/internal/querynodev2/server.go @@ -269,6 +269,12 @@ func (node *QueryNode) InitSegcore() error { return err } + tieredStorageEnabledGlobally := C.bool(paramtable.Get().QueryNodeCfg.TieredStorageEnableGlobal.GetAsBool()) + tieredStorageMemoryLimit := C.int64_t(paramtable.Get().QueryNodeCfg.TieredStorageMemoryAllocationRatio.GetAsFloat() * float64(hardware.GetMemoryCount())) + diskCapacity := paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64() + tieredStorageDiskLimit := C.int64_t(paramtable.Get().QueryNodeCfg.TieredStorageDiskAllocationRatio.GetAsFloat() * float64(diskCapacity)) + C.ConfigureTieredStorage(tieredStorageEnabledGlobally, tieredStorageMemoryLimit, tieredStorageDiskLimit) + initcore.InitTraceConfig(paramtable.Get()) C.InitExecExpressionFunctionFactory() return nil diff --git a/internal/util/segcore/segment.go b/internal/util/segcore/segment.go index ffa1ee8945..00d71445e4 100644 --- a/internal/util/segcore/segment.go +++ b/internal/util/segcore/segment.go @@ -229,8 +229,6 @@ func (s *cSegmentImpl) preInsert(numOfRecords int) (int64, error) { // Delete deletes entities from the segment. func (s *cSegmentImpl) Delete(ctx context.Context, request *DeleteRequest) (*DeleteResult, error) { - cOffset := C.int64_t(0) // depre - cSize := C.int64_t(request.PrimaryKeys.Len()) cTimestampsPtr := (*C.uint64_t)(&(request.Timestamps)[0]) @@ -244,7 +242,6 @@ func (s *cSegmentImpl) Delete(ctx context.Context, request *DeleteRequest) (*Del return nil, fmt.Errorf("failed to marshal ids: %s", err) } status := C.Delete(s.ptr, - cOffset, cSize, (*C.uint8_t)(unsafe.Pointer(&dataBlob[0])), (C.uint64_t)(len(dataBlob)), diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 4b28d16161..42f863135a 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -2686,6 +2686,10 @@ type queryNodeConfig struct { InterimIndexMemExpandRate ParamItem `refreshable:"false"` InterimIndexBuildParallelRate ParamItem `refreshable:"false"` MultipleChunkedEnable ParamItem `refreshable:"false"` // Deprecated + // TODO this should be refreshable? + TieredStorageEnableGlobal ParamItem `refreshable:"false"` + TieredStorageMemoryAllocationRatio ParamItem `refreshable:"false"` + TieredStorageDiskAllocationRatio ParamItem `refreshable:"false"` KnowhereScoreConsistency ParamItem `refreshable:"false"` @@ -2823,6 +2827,47 @@ func (p *queryNodeConfig) init(base *BaseTable) { } p.StatsPublishInterval.Init(base.mgr) + p.TieredStorageEnableGlobal = ParamItem{ + Key: "queryNode.segcore.tieredStorage.enableGlobally", + Version: "2.6.0", + DefaultValue: "false", + Doc: "Whether or not to turn on Tiered Storage globally in this cluster.", + Export: true, + } + p.TieredStorageEnableGlobal.Init(base.mgr) + + p.TieredStorageMemoryAllocationRatio = ParamItem{ + Key: "queryNode.segcore.tieredStorage.memoryAllocationRatio", + Version: "2.6.0", + DefaultValue: "0.5", + Formatter: func(v string) string { + ratio := getAsFloat(v) + if ratio < 0 || ratio > 1 { + return "0.5" + } + return fmt.Sprintf("%f", ratio) + }, + Doc: "The ratio of memory allocation for Tiered Storage.", + Export: true, + } + p.TieredStorageMemoryAllocationRatio.Init(base.mgr) + + p.TieredStorageDiskAllocationRatio = ParamItem{ + Key: "queryNode.segcore.tieredStorage.diskAllocationRatio", + Version: "2.6.0", + DefaultValue: "0.5", + Formatter: func(v string) string { + ratio := getAsFloat(v) + if ratio < 0 || ratio > 1 { + return "0.5" + } + return fmt.Sprintf("%f", ratio) + }, + Doc: "The ratio of disk allocation for Tiered Storage.", + Export: true, + } + p.TieredStorageDiskAllocationRatio.Init(base.mgr) + p.KnowhereThreadPoolSize = ParamItem{ Key: "queryNode.segcore.knowhereThreadPoolNumRatio", Version: "2.0.0", diff --git a/scripts/run_cpp_unittest.sh b/scripts/run_cpp_unittest.sh index 636fddcd30..110df7148c 100755 --- a/scripts/run_cpp_unittest.sh +++ b/scripts/run_cpp_unittest.sh @@ -71,6 +71,13 @@ for UNITTEST_DIR in "${UNITTEST_DIRS[@]}"; do fi fi + echo "Running cachinglayer unittest" + ${UNITTEST_DIR}/cachinglayer_test + if [ $? -ne 0 ]; then + echo ${UNITTEST_DIR}/cachinglayer_test "run failed" + exit 1 + fi + done # run cwrapper unittest