diff --git a/examples/cpp/406_feature_read_cache.cpp b/examples/cpp/406_feature_read_cache.cpp new file mode 100644 index 0000000000..b13cb0bb19 --- /dev/null +++ b/examples/cpp/406_feature_read_cache.cpp @@ -0,0 +1,144 @@ +// Copyright 2024-present the vsag project +// +// 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 + +#include +#include +#include +#include +#include + +int +main(int argc, char** argv) { + vsag::init(); + + /******************* Prepare Base Dataset *****************/ + int64_t num_vectors = 1000; + int64_t dim = 64; + std::vector ids(num_vectors); + std::vector datas(num_vectors * dim); + std::mt19937 rng(47); + std::uniform_real_distribution distrib_real; + for (int64_t i = 0; i < num_vectors; ++i) { + ids[i] = i; + } + for (int64_t i = 0; i < dim * num_vectors; ++i) { + datas[i] = distrib_real(rng); + } + auto base = vsag::Dataset::Make(); + base->NumElements(num_vectors) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(datas.data()) + ->Owner(false); + + std::string precise_file_path = argc > 1 ? argv[1] : "vsag_read_cache_precise"; + + /******************* Create HGraph Index with ReadCache on precise codes *****************/ + std::string hgraph_build_parameters = R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 64, + "index_param": { + "base_quantization_type": "sq8", + "max_degree": 16, + "ef_construction": 40, + "alpha": 1.2, + "use_reorder": true, + "precise_quantization_type": "fp32", + "base_io_type": "memory_io", + "precise_io_type": "async_io", + "precise_enable_read_cache": true, + "precise_file_path": ")" + precise_file_path + + R"(", + "precise_cache_total_size": 131072 + } + } + )"; + vsag::Resource resource(vsag::Engine::CreateDefaultAllocator(), nullptr); + vsag::Engine engine(&resource); + vsag::Options::Instance().set_block_size_limit(2 * 1024 * 1024); + auto index_result = engine.CreateIndex("hgraph", hgraph_build_parameters); + if (not index_result.has_value()) { + std::cerr << "Failed to create index: " << index_result.error().message << std::endl; + return -1; + } + auto index = index_result.value(); + + /******************* Build HGraph Index *****************/ + if (auto build_result = index->Build(base); build_result.has_value()) { + std::cout << "After Build(), Index HGraph with ReadCache contains: " + << index->GetNumElements() << std::endl; + } else { + std::cerr << "Failed to build index: " << build_result.error().message << std::endl; + return -1; + } + + /******************* Serialize to Disk *****************/ + auto serialize_result = index->Serialize(); + if (not serialize_result.has_value()) { + auto error = serialize_result.error(); + std::cerr << "Failed to serialize index: " << error.message << std::endl; + return -1; + } + + /******************* Deserialize from Disk *****************/ + auto deserialize_result = vsag::Factory::CreateIndex("hgraph", hgraph_build_parameters); + if (not deserialize_result.has_value()) { + std::cerr << "Failed to create index for deserialization: " + << deserialize_result.error().message << std::endl; + return -1; + } + auto restored_index = deserialize_result.value(); + if (auto ds_result = restored_index->Deserialize(serialize_result.value()); + not ds_result.has_value()) { + std::cerr << "Failed to deserialize index: " << ds_result.error().message << std::endl; + return -1; + } + + /******************* Prepare Query Dataset *****************/ + std::vector query_vector(dim); + for (int64_t i = 0; i < dim; ++i) { + query_vector[i] = distrib_real(rng); + } + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector.data())->Owner(false); + + /******************* KnnSearch For Restored Index *****************/ + auto hgraph_search_parameters = R"( + { + "hgraph": { + "ef_search": 40 + } + } + )"; + int64_t topk = 10; + auto search_result = restored_index->KnnSearch(query, topk, hgraph_search_parameters); + if (not search_result.has_value()) { + std::cerr << "Failed to search index: " << search_result.error().message << std::endl; + return -1; + } + auto result = search_result.value(); + + /******************* Print Search Result *****************/ + std::cout << "results: " << std::endl; + for (int64_t i = 0; i < result->GetDim(); ++i) { + std::cout << result->GetIds()[i] << ": " << result->GetDistances()[i] << std::endl; + } + + engine.Shutdown(); + return 0; +} diff --git a/examples/cpp/407_feature_ivf_read_cache.cpp b/examples/cpp/407_feature_ivf_read_cache.cpp new file mode 100644 index 0000000000..96bde69e53 --- /dev/null +++ b/examples/cpp/407_feature_ivf_read_cache.cpp @@ -0,0 +1,140 @@ +// Copyright 2024-present the vsag project +// +// 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 + +#include +#include +#include +#include +#include + +int +main(int argc, char** argv) { + vsag::init(); + + /******************* Prepare Base Dataset *****************/ + int64_t num_vectors = 1000; + int64_t dim = 64; + std::vector ids(num_vectors); + std::vector datas(num_vectors * dim); + std::mt19937 rng(47); + std::uniform_real_distribution distrib_real; + for (int64_t i = 0; i < num_vectors; ++i) { + ids[i] = i; + } + for (int64_t i = 0; i < dim * num_vectors; ++i) { + datas[i] = distrib_real(rng); + } + auto base = vsag::Dataset::Make(); + base->NumElements(num_vectors) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(datas.data()) + ->Owner(false); + + std::string precise_file_path = argc > 1 ? argv[1] : "vsag_ivf_read_cache_precise"; + + /******************* Create IVF Index with ReadCache on reorder codes *****************/ + std::string ivf_build_params = R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 64, + "index_param": { + "buckets_count": 10, + "base_quantization_type": "fp32", + "partition_strategy_type": "ivf", + "ivf_train_type": "kmeans", + "train_sample_count": 800, + "use_reorder": true, + "precise_quantization_type": "fp32", + "base_io_type": "memory_io", + "precise_io_type": "async_io", + "precise_enable_read_cache": true, + "precise_file_path": ")" + + precise_file_path + R"(", + "precise_cache_total_size": 131072 + } + } + )"; + auto index_result = vsag::Factory::CreateIndex("ivf", ivf_build_params); + if (not index_result.has_value()) { + std::cerr << "Failed to create index: " << index_result.error().message << std::endl; + return -1; + } + auto index = index_result.value(); + + /******************* Build IVF Index *****************/ + if (auto build_result = index->Build(base); build_result.has_value()) { + std::cout << "After Build(), Index IVF with reorder ReadCache contains: " + << index->GetNumElements() << std::endl; + } else { + std::cerr << "Failed to build index: " << build_result.error().message << std::endl; + return -1; + } + + /******************* Serialize to Disk *****************/ + auto serialize_result = index->Serialize(); + if (not serialize_result.has_value()) { + auto error = serialize_result.error(); + std::cerr << "Failed to serialize index: " << error.message << std::endl; + return -1; + } + + /******************* Deserialize from Disk *****************/ + auto deserialize_result = vsag::Factory::CreateIndex("ivf", ivf_build_params); + if (not deserialize_result.has_value()) { + std::cerr << "Failed to create index for deserialization: " + << deserialize_result.error().message << std::endl; + return -1; + } + auto restored_index = deserialize_result.value(); + if (auto ds_result = restored_index->Deserialize(serialize_result.value()); + not ds_result.has_value()) { + std::cerr << "Failed to deserialize index: " << ds_result.error().message << std::endl; + return -1; + } + + /******************* Prepare Query Dataset *****************/ + std::vector query_vector(dim); + for (int64_t i = 0; i < dim; ++i) { + query_vector[i] = distrib_real(rng); + } + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector.data())->Owner(false); + + /******************* KnnSearch For Restored Index *****************/ + auto ivf_search_parameters = R"( + { + "ivf": { + "scan_buckets_count": 4 + } + })"; + int64_t topk = 10; + auto search_result = restored_index->KnnSearch(query, topk, ivf_search_parameters); + if (not search_result.has_value()) { + std::cerr << "Failed to search index: " << search_result.error().message << std::endl; + return -1; + } + auto result = search_result.value(); + + /******************* Print Search Result *****************/ + std::cout << "results: " << std::endl; + for (int64_t i = 0; i < result->GetDim(); ++i) { + std::cout << result->GetIds()[i] << ": " << result->GetDistances()[i] << std::endl; + } + + return 0; +} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 3793625a87..899b1228e0 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -154,6 +154,11 @@ target_link_libraries(324_feature_lazy_hgraph_extra_info vsag) add_executable(325_feature_uring_io 325_feature_uring_io.cpp) target_link_libraries(325_feature_uring_io vsag) - add_executable(324_feature_hgraph_mci_companion 324_feature_hgraph_mci_companion.cpp) target_link_libraries(324_feature_hgraph_mci_companion vsag) + +add_executable(406_feature_read_cache 406_feature_read_cache.cpp) +target_link_libraries(406_feature_read_cache vsag) + +add_executable(407_feature_ivf_read_cache 407_feature_ivf_read_cache.cpp) +target_link_libraries(407_feature_ivf_read_cache vsag) diff --git a/include/vsag/constants.h b/include/vsag/constants.h index ceb316c1fb..ed81351b78 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -219,6 +219,14 @@ extern const char* const HGRAPH_BASE_PQ_DIM; extern const char* const HGRAPH_BASE_FILE_PATH; extern const char* const HGRAPH_BASE_DIRECT_READ; extern const char* const HGRAPH_BASE_SUPPLEMENT_FILE_PATH; +extern const char* const HGRAPH_BASE_ENABLE_READ_CACHE; +extern const char* const HGRAPH_BASE_CACHE_TOTAL_SIZE; +extern const char* const HGRAPH_GRAPH_ENABLE_READ_CACHE; +extern const char* const HGRAPH_GRAPH_CACHE_TOTAL_SIZE; +extern const char* const HGRAPH_PRECISE_ENABLE_READ_CACHE; +extern const char* const HGRAPH_PRECISE_CACHE_TOTAL_SIZE; +extern const char* const HGRAPH_RAW_VECTOR_ENABLE_READ_CACHE; +extern const char* const HGRAPH_RAW_VECTOR_CACHE_TOTAL_SIZE; extern const char* const HGRAPH_PRECISE_IO_TYPE; extern const char* const HGRAPH_PRECISE_FILE_PATH; extern const char* const HGRAPH_PRECISE_DIRECT_READ; @@ -261,6 +269,10 @@ extern const char* const IVF_BASE_QUANTIZATION_TYPE; extern const char* const IVF_BASE_IO_TYPE; extern const char* const IVF_BASE_PQ_DIM; extern const char* const IVF_BASE_FILE_PATH; +extern const char* const IVF_BASE_ENABLE_READ_CACHE; +extern const char* const IVF_BASE_CACHE_TOTAL_SIZE; +extern const char* const IVF_PRECISE_ENABLE_READ_CACHE; +extern const char* const IVF_PRECISE_CACHE_TOTAL_SIZE; extern const char* const IVF_PRECISE_QUANTIZATION_TYPE; extern const char* const IVF_PRECISE_IO_TYPE; extern const char* const IVF_PRECISE_FILE_PATH; diff --git a/src/algorithm/hgraph/hgraph_param_mapping.cpp b/src/algorithm/hgraph/hgraph_param_mapping.cpp index 073e9650bc..7acbdb8e05 100644 --- a/src/algorithm/hgraph/hgraph_param_mapping.cpp +++ b/src/algorithm/hgraph/hgraph_param_mapping.cpp @@ -193,6 +193,7 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { IO_FILE_PATH_KEY, }, }, + // clang-format off { HGRAPH_BASE_DIRECT_READ, { @@ -201,6 +202,14 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { IO_DIRECT_READ_KEY, }, }, + { + HGRAPH_BASE_CACHE_TOTAL_SIZE, + { + BASE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, + }, { HGRAPH_PRECISE_FILE_PATH, { @@ -217,6 +226,15 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { IO_DIRECT_READ_KEY, }, }, + { + HGRAPH_PRECISE_CACHE_TOTAL_SIZE, + { + PRECISE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, + }, + // clang-format on { HGRAPH_PRECISE_QUANTIZATION_TYPE, { @@ -241,6 +259,14 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { IO_FILE_PATH_KEY, }, }, + { + HGRAPH_GRAPH_CACHE_TOTAL_SIZE, + { + GRAPH_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, + }, { STORE_RAW_VECTOR, { @@ -562,6 +588,46 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { { HGRAPH_MCI_INCREMENTAL_CLIQUE_MAX_KEY, }, + }, + { + HGRAPH_BASE_ENABLE_READ_CACHE, + { + BASE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, + { + HGRAPH_PRECISE_ENABLE_READ_CACHE, + { + PRECISE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, + { + HGRAPH_GRAPH_ENABLE_READ_CACHE, + { + GRAPH_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, + { + HGRAPH_RAW_VECTOR_ENABLE_READ_CACHE, + { + RAW_VECTOR_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, + { + HGRAPH_RAW_VECTOR_CACHE_TOTAL_SIZE, + { + RAW_VECTOR_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, }}; const std::string hgraph_params_template = R"( diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index 9a3f4d31e1..46c45dd838 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -140,6 +140,14 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param, IO_FILE_PATH_KEY, }, }, + { + IVF_BASE_CACHE_TOTAL_SIZE, + { + BUCKET_PARAMS_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, + }, { IVF_PRECISE_QUANTIZATION_TYPE, { @@ -164,6 +172,14 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param, IO_FILE_PATH_KEY, }, }, + { + IVF_PRECISE_CACHE_TOTAL_SIZE, + { + PRECISE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_TOTAL_CACHE_SIZE_KEY, + }, + }, { IVF_BUCKETS_COUNT, { @@ -332,6 +348,22 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param, GRAPH_BUILD_THRESHOLD_KEY, }, }, + { + IVF_BASE_ENABLE_READ_CACHE, + { + BUCKET_PARAMS_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, + { + IVF_PRECISE_ENABLE_READ_CACHE, + { + PRECISE_CODES_KEY, + IO_PARAMS_KEY, + READ_CACHE_ENABLED_KEY, + }, + }, }; if (common_param.data_type_ == DataTypes::DATA_TYPE_INT8) { diff --git a/src/constants.cpp b/src/constants.cpp index 9dd4a50628..4a0c842d85 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -188,6 +188,14 @@ const char* const HGRAPH_BASE_PQ_DIM = "base_pq_dim"; const char* const HGRAPH_BASE_FILE_PATH = "base_file_path"; const char* const HGRAPH_BASE_DIRECT_READ = "base_direct_read"; const char* const HGRAPH_BASE_SUPPLEMENT_FILE_PATH = "base_supplement_file_path"; +const char* const HGRAPH_BASE_ENABLE_READ_CACHE = "base_enable_read_cache"; +const char* const HGRAPH_BASE_CACHE_TOTAL_SIZE = "base_cache_total_size"; +const char* const HGRAPH_GRAPH_ENABLE_READ_CACHE = "graph_enable_read_cache"; +const char* const HGRAPH_GRAPH_CACHE_TOTAL_SIZE = "graph_cache_total_size"; +const char* const HGRAPH_PRECISE_ENABLE_READ_CACHE = "precise_enable_read_cache"; +const char* const HGRAPH_PRECISE_CACHE_TOTAL_SIZE = "precise_cache_total_size"; +const char* const HGRAPH_RAW_VECTOR_ENABLE_READ_CACHE = "raw_vector_enable_read_cache"; +const char* const HGRAPH_RAW_VECTOR_CACHE_TOTAL_SIZE = "raw_vector_cache_total_size"; const char* const HGRAPH_PRECISE_IO_TYPE = "precise_io_type"; const char* const HGRAPH_PRECISE_FILE_PATH = "precise_file_path"; const char* const HGRAPH_PRECISE_DIRECT_READ = "precise_direct_read"; @@ -230,6 +238,10 @@ const char* const IVF_BASE_QUANTIZATION_TYPE = "base_quantization_type"; const char* const IVF_BASE_IO_TYPE = "base_io_type"; const char* const IVF_BASE_PQ_DIM = "base_pq_dim"; const char* const IVF_BASE_FILE_PATH = "base_file_path"; +const char* const IVF_BASE_ENABLE_READ_CACHE = "base_enable_read_cache"; +const char* const IVF_BASE_CACHE_TOTAL_SIZE = "base_cache_total_size"; +const char* const IVF_PRECISE_ENABLE_READ_CACHE = "precise_enable_read_cache"; +const char* const IVF_PRECISE_CACHE_TOTAL_SIZE = "precise_cache_total_size"; const char* const PYRAMID_SUPPORT_DUPLICATE = SUPPORT_DUPLICATE; const char* const PYRAMID_EF_CONSTRUCTION = EF_CONSTRUCTION_KEY; diff --git a/src/datacell/bucket_datacell.h b/src/datacell/bucket_datacell.h index 57ca26a43a..dcdc0b5b80 100644 --- a/src/datacell/bucket_datacell.h +++ b/src/datacell/bucket_datacell.h @@ -23,12 +23,25 @@ #include "bucket_interface.h" #include "impl/inner_search_param.h" #include "io/container/io_array.h" +#include "io/read_cache/page.h" #include "quantization/product_quantization/pq_fastscan_quantizer.h" #include "simd/fp32_simd.h" #include "utils/byte_buffer.h" namespace vsag { +inline IOParamPtr +AdjustBucketReadCacheParam(const IOParamPtr& io_param, BucketIdType bucket_count) { + if (io_param == nullptr or not io_param->enable_read_cache_ or bucket_count == 0) { + return io_param; + } + JsonType json = io_param->ToJson(); + uint64_t total_pages = io_param->read_cache_total_size_ / Page::DEFAULT_PAGE_SIZE; + uint64_t pages_per_bucket = total_pages / bucket_count; + json[READ_CACHE_TOTAL_CACHE_SIZE_KEY].SetUint64(pages_per_bucket * Page::DEFAULT_PAGE_SIZE); + return IOParameter::GetIOParameterByJson(json); +} + template class BucketDataCell : public BucketInterface { public: @@ -203,7 +216,9 @@ BucketDataCell::BucketDataCell(const QuantizerParamPtr& quant BucketIdType bucket_count, bool use_residual) : BucketInterface(), - datas_(common_param.allocator_.get(), io_param, common_param), + datas_(common_param.allocator_.get(), + AdjustBucketReadCacheParam(io_param, bucket_count), + common_param), bucket_sizes_(bucket_count, 0, common_param.allocator_.get()), inner_ids_(bucket_count, Vector(common_param.allocator_.get()), diff --git a/src/datacell/bucket_interface.cpp b/src/datacell/bucket_interface.cpp index bcf7abef30..58d36c6d74 100644 --- a/src/datacell/bucket_interface.cpp +++ b/src/datacell/bucket_interface.cpp @@ -15,7 +15,9 @@ #include "bucket_interface.h" #include "bucket_interface_factory.h" +#include "bucket_interface_factory_impl.h" #include "inner_string_params.h" +#include "io/io_headers.h" namespace vsag { diff --git a/src/datacell/rabitq_split_datacell.h b/src/datacell/rabitq_split_datacell.h index f6cddb1fd2..18fbe05a44 100644 --- a/src/datacell/rabitq_split_datacell.h +++ b/src/datacell/rabitq_split_datacell.h @@ -179,10 +179,11 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil // memory + supplement on disk). Otherwise fall back to the shared // io_param with the legacy file-path suffix to keep the two backing // files separate for file-backed IO. - const IOParamPtr one_bit_io_param = SuffixIOParam(io_param, "_onebit"); + const bool shares_io_param = supplement_io_param == nullptr; + const IOParamPtr one_bit_io_param = SuffixIOParam(io_param, "_onebit", shares_io_param); const IOParamPtr supp_io_param = (supplement_io_param != nullptr) ? supplement_io_param - : SuffixIOParam(io_param, "_supplement"); + : SuffixIOParam(io_param, "_supplement", true); if (supplement_io_param != nullptr) { this->supplement_io_type_ = supplement_io_param->GetTypeName(); } @@ -753,7 +754,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void InitIO(const IOParamPtr& io_param) override { - this->x_bit_cell_->InitIO(SuffixIOParam(io_param, "_onebit")); + const bool shares_io_param = this->supplement_io_type_.empty(); + this->x_bit_cell_->InitIO(SuffixIOParam(io_param, "_onebit", shares_io_param)); // In hybrid mode (one-bit and supplement use different IO backends) // the caller-facing `io_param` is the one-bit IO parameter type and // cannot be passed directly to `supplement_cell_`. Rebuild a fresh @@ -764,7 +766,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void InitIO(const IOParamPtr& one_bit_io_param, const IOParamPtr& supplement_io_param) { - this->x_bit_cell_->InitIO(SuffixIOParam(one_bit_io_param, "_onebit")); + const bool shares_io_param = supplement_io_param == nullptr; + this->x_bit_cell_->InitIO(SuffixIOParam(one_bit_io_param, "_onebit", shares_io_param)); if (supplement_io_param != nullptr) { // Refresh the recorded supplement type so subsequent // single-parameter InitIO calls (e.g. from Deserialize) can @@ -972,7 +975,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil private: static IOParamPtr - SuffixIOParam(const IOParamPtr& io_param, const std::string& suffix) { + SuffixIOParam(const IOParamPtr& io_param, const std::string& suffix, bool split_cache = false) { if (io_param == nullptr) { return nullptr; } @@ -981,6 +984,9 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil std::string path = json[IO_FILE_PATH_KEY].GetString(); json[IO_FILE_PATH_KEY].SetString(path + suffix); } + if (split_cache and io_param->enable_read_cache_) { + json[READ_CACHE_TOTAL_CACHE_SIZE_KEY].SetUint64(io_param->read_cache_total_size_ / 2); + } return IOParameter::GetIOParameterByJson(json); } @@ -998,7 +1004,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil return nullptr; } if (this->supplement_io_type_.empty()) { - return SuffixIOParam(io_param, "_supplement"); + return SuffixIOParam(io_param, "_supplement", true); } auto json = io_param->ToJson(); json[TYPE_KEY].SetString(this->supplement_io_type_); diff --git a/src/inner_string_params.h b/src/inner_string_params.h index fbbf7e9312..6c2115a0f8 100644 --- a/src/inner_string_params.h +++ b/src/inner_string_params.h @@ -71,6 +71,8 @@ const char* const IO_TYPE_VALUE_READER_IO = "reader_io"; const char* const IO_TYPE_VALUE_ASYNC_IO = "async_io"; const char* const IO_TYPE_VALUE_URING_IO = "uring_io"; const char* const IO_TYPE_VALUE_BLOCK_MEMORY_IO = "block_memory_io"; +const char* const READ_CACHE_TOTAL_CACHE_SIZE_KEY = "total_cache_size"; +const char* const READ_CACHE_ENABLED_KEY = "enable_read_cache"; const char* const BLOCK_IO_BLOCK_SIZE_KEY = "block_size"; // IO param for file @@ -229,6 +231,8 @@ const std::unordered_map DEFAULT_MAP = { {"IO_TYPE_VALUE_MEMORY_IO", IO_TYPE_VALUE_MEMORY_IO}, {"IO_TYPE_VALUE_BLOCK_MEMORY_IO", IO_TYPE_VALUE_BLOCK_MEMORY_IO}, {"IO_TYPE_VALUE_BUFFER_IO", IO_TYPE_VALUE_BUFFER_IO}, + {"READ_CACHE_TOTAL_CACHE_SIZE_KEY", READ_CACHE_TOTAL_CACHE_SIZE_KEY}, + {"READ_CACHE_ENABLED_KEY", READ_CACHE_ENABLED_KEY}, {"IO_PARAMS_KEY", IO_PARAMS_KEY}, {"BLOCK_IO_BLOCK_SIZE_KEY", BLOCK_IO_BLOCK_SIZE_KEY}, {"QUANTIZATION_TYPE_VALUE_SQ8", QUANTIZATION_TYPE_VALUE_SQ8}, diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index d9487149b3..9213a9b71e 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -33,6 +33,8 @@ set (IO_SRC memory_block_io/memory_block_io.cpp reader_io/reader_io.cpp reader_io/reader_io_parameter.cpp + read_cache/page_cache.cpp + read_cache/lru_page_cache.cpp ) add_library (io OBJECT ${IO_SRC}) diff --git a/src/io/async_io/async_io.cpp b/src/io/async_io/async_io.cpp index 322ce5a1f6..4e240e7dfe 100644 --- a/src/io/async_io/async_io.cpp +++ b/src/io/async_io/async_io.cpp @@ -54,7 +54,9 @@ AsyncIO::AsyncIO(const AsyncIOParameterPtr& io_param, const IndexCommonParam& co : AsyncIO(io_param->path_, common_param.allocator_.get()){}; AsyncIO::AsyncIO(const IOParamPtr& param, const IndexCommonParam& common_param) - : AsyncIO(std::dynamic_pointer_cast(param), common_param){}; + : AsyncIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); +}; AsyncIO::~AsyncIO() { close(this->wfd_); diff --git a/src/io/async_io/async_io_parameter.cpp b/src/io/async_io/async_io_parameter.cpp index 19e832c700..99d400831e 100644 --- a/src/io/async_io/async_io_parameter.cpp +++ b/src/io/async_io/async_io_parameter.cpp @@ -37,6 +37,7 @@ AsyncIOParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_ASYNC_IO); json[IO_FILE_PATH_KEY].SetString(this->path_); + AppendReadCacheConfig(json); return json; } } // namespace vsag diff --git a/src/io/buffer_io/buffer_io.cpp b/src/io/buffer_io/buffer_io.cpp index eb950dae43..1a915b8c55 100644 --- a/src/io/buffer_io/buffer_io.cpp +++ b/src/io/buffer_io/buffer_io.cpp @@ -44,6 +44,7 @@ BufferIO::BufferIO(const BufferIOParameterPtr& io_param, const IndexCommonParam& BufferIO::BufferIO(const IOParamPtr& param, const IndexCommonParam& common_param) : BufferIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); } void diff --git a/src/io/buffer_io/buffer_io_parameter.cpp b/src/io/buffer_io/buffer_io_parameter.cpp index 883826c139..7cff0536e8 100644 --- a/src/io/buffer_io/buffer_io_parameter.cpp +++ b/src/io/buffer_io/buffer_io_parameter.cpp @@ -38,6 +38,7 @@ BufferIOParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_BUFFER_IO); json[IO_FILE_PATH_KEY].SetString(this->path_); + AppendReadCacheConfig(json); return json; } } // namespace vsag diff --git a/src/io/common/basic_io.h b/src/io/common/basic_io.h index b161e69780..94cd6080a5 100644 --- a/src/io/common/basic_io.h +++ b/src/io/common/basic_io.h @@ -20,8 +20,14 @@ #include #include #include +#include +#include +#include #include "io/common/io_parameter.h" +#include "io/read_cache/lru_page_cache.h" +#include "io/read_cache/page.h" +#include "io/read_cache/page_cache.h" #include "storage/stream_reader.h" #include "storage/stream_writer.h" #include "utils/byte_buffer.h" @@ -29,6 +35,13 @@ namespace vsag { +template +struct SupportsZeroSizeResize : std::true_type {}; + +template +struct SupportsZeroSizeResize> + : std::bool_constant {}; + /** * @brief A template class for basic input/output operations. * @@ -68,6 +81,9 @@ class BasicIO { Write(const uint8_t* data, uint64_t size, uint64_t offset) { static_assert(has_WriteImpl::value); cast().WriteImpl(data, size, offset); + if constexpr (not InMemory) { + InvalidateCacheRange(size, offset); + } } /** @@ -84,6 +100,11 @@ class BasicIO { inline bool Read(uint64_t size, uint64_t offset, uint8_t* data) const { static_assert(has_ReadImpl::value); + if constexpr (not InMemory) { + if (cache_ != nullptr) { + return ReadCached(size, offset, data); + } + } return cast().ReadImpl(size, offset, data); } @@ -101,6 +122,24 @@ class BasicIO { [[nodiscard]] inline const uint8_t* Read(uint64_t size, uint64_t offset, bool& need_release) const { static_assert(has_DirectReadImpl::value); + if constexpr (not InMemory) { + if (cache_ != nullptr) { + need_release = false; + if (size == 0 or not IsValidRange(size, offset)) { + return nullptr; + } + auto* data = static_cast(allocator_->Allocate(size)); + if (data == nullptr) { + return nullptr; + } + if (not ReadCached(size, offset, data)) { + allocator_->Deallocate(data); + return nullptr; + } + need_release = true; + return data; + } + } return cast().DirectReadImpl(size, offset, need_release); // TODO(LHT129): use IOReadObject } @@ -119,6 +158,17 @@ class BasicIO { inline bool MultiRead(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { static_assert(has_MultiReadImpl::value); + if constexpr (not InMemory) { + if (cache_ != nullptr) { + for (uint64_t i = 0; i < count; ++i) { + if (not ReadCached(sizes[i], offsets[i], datas)) { + return false; + } + datas += sizes[i]; + } + return true; + } + } return cast().MultiReadImpl(datas, sizes, offsets, count); } @@ -165,11 +215,25 @@ class BasicIO { Deserialize(StreamReader& reader) { uint64_t size = 0; StreamReader::ReadObj(reader, size); + has_deserialized_ = true; this->start_ = reader.GetCursor(); if constexpr (SkipDeserialize) { reader.Seek(reader.GetCursor() + size); this->size_ = std::max(this->size_, size); } else { + // Reset the logical and physical extent so a shorter deserialization + // cannot retain stale bytes from a previously opened file. + if constexpr (has_ResizeImpl::value) { + if (size > 0 or SupportsZeroSizeResize::value) { + Resize(size); + } else { + this->size_ = 0; + ClearCache(); + } + } else { + this->size_ = 0; + ClearCache(); + } ByteBuffer buffer(SERIALIZE_BUFFER_SIZE, this->allocator_); uint64_t offset = 0; while (offset < size) { @@ -194,6 +258,12 @@ class BasicIO { inline void Release(const uint8_t* data) const { + if constexpr (not InMemory) { + if (cache_ != nullptr) { + allocator_->Deallocate(const_cast(data)); + return; + } + } if constexpr (has_ReleaseImpl::value) { return cast().ReleaseImpl(data); } @@ -210,6 +280,15 @@ class BasicIO { */ inline void InitIO(const IOParamPtr& io_param) { + if constexpr (not InMemory) { + if (cache_ == nullptr) { + EnableReadCache(io_param); + } else if (io_param != nullptr and io_param->enable_read_cache_) { + EnableReadCache(io_param); + } else { + ClearCache(); + } + } if constexpr (has_InitIOImpl::value) { return cast().InitIOImpl(io_param); } @@ -218,7 +297,7 @@ class BasicIO { inline void Resize(uint64_t size) { if constexpr (has_ResizeImpl::value) { - return cast().ResizeImpl(size); + cast().ResizeImpl(size); } else { if (size <= this->size_) { return; @@ -232,17 +311,23 @@ class BasicIO { offset += cur_size; } } + if constexpr (not InMemory) { + ClearCache(); + } } inline void Shrink(uint64_t size) { if constexpr (has_ShrinkImpl::value) { - return cast().ShrinkImpl(size); + cast().ShrinkImpl(size); } else { if (size <= this->size_) { this->size_ = size; } } + if constexpr (not InMemory) { + ClearCache(); + } } inline int64_t @@ -253,6 +338,27 @@ class BasicIO { return this->size_; } + [[nodiscard]] bool + HasDeserialized() const { + return has_deserialized_; + } + + void + EnableReadCache(const IOParamPtr& io_param) { + if constexpr (not InMemory) { + if (io_param == nullptr or not io_param->enable_read_cache_) { + cache_.reset(); + return; + } + auto page_count = io_param->read_cache_total_size_ / Page::DEFAULT_PAGE_SIZE; + if (page_count == 0) { + cache_.reset(); + return; + } + cache_ = std::make_unique(page_count); + } + } + public: /** * @brief The size of the IO object. @@ -320,11 +426,91 @@ class BasicIO { return static_cast(*this); } + [[nodiscard]] bool + IsValidRange(uint64_t size, uint64_t offset) const { + return offset <= size_ and size <= size_ - offset; + } + + bool + ReadCached(uint64_t size, uint64_t offset, uint8_t* data) const { + if (not IsValidRange(size, offset)) { + return false; + } + uint64_t copied = 0; + while (copied < size) { + uint64_t current_offset = offset + copied; + uint64_t page_id = current_offset / Page::DEFAULT_PAGE_SIZE; + uint64_t page_offset = current_offset % Page::DEFAULT_PAGE_SIZE; + uint64_t copy_size = std::min(size - copied, Page::DEFAULT_PAGE_SIZE - page_offset); + auto page = GetOrLoadPage(page_id); + if (page == nullptr) { + return false; + } + std::memcpy(data + copied, page->Data() + page_offset, copy_size); + copied += copy_size; + } + return true; + } + + PagePtr + GetOrLoadPage(uint64_t page_id) const { + if (page_id > UINT64_MAX / Page::DEFAULT_PAGE_SIZE) { + return nullptr; + } + uint64_t offset = page_id * Page::DEFAULT_PAGE_SIZE; + if (offset >= size_) { + return nullptr; + } + std::scoped_lock lock(cache_mutex_); + auto page = cache_->Get(page_id); + if (page != nullptr) { + return page; + } + auto new_page = std::make_shared(allocator_); + if (new_page->Data() == nullptr) { + return nullptr; + } + uint64_t read_size = std::min(Page::DEFAULT_PAGE_SIZE, size_ - offset); + if (not cast().ReadImpl(read_size, offset, new_page->Data())) { + return nullptr; + } + return cache_->Insert(page_id, std::move(new_page)); + } + + void + InvalidateCacheRange(uint64_t size, uint64_t offset) { + if (cache_ == nullptr or size == 0) { + return; + } + std::scoped_lock lock(cache_mutex_); + if (offset > UINT64_MAX - (size - 1)) { + cache_->Clear(); + return; + } + uint64_t first_page = offset / Page::DEFAULT_PAGE_SIZE; + uint64_t last_page = (offset + size - 1) / Page::DEFAULT_PAGE_SIZE; + for (uint64_t page_id = first_page; page_id <= last_page; ++page_id) { + cache_->Remove(page_id); + } + } + + void + ClearCache() { + if (cache_ != nullptr) { + std::scoped_lock lock(cache_mutex_); + cache_->Clear(); + } + } + /** * @brief The size of the max buffer used for serialization. */ static constexpr uint64_t SERIALIZE_BUFFER_SIZE = 1024 * 1024 * 2; + mutable std::mutex cache_mutex_; + mutable std::unique_ptr cache_; + bool has_deserialized_{false}; + private: /** * @brief Generates a struct to check if a class has a member function with a specific signature. diff --git a/src/io/common/io_parameter.cpp b/src/io/common/io_parameter.cpp index 2b1200be9e..aba2461976 100644 --- a/src/io/common/io_parameter.cpp +++ b/src/io/common/io_parameter.cpp @@ -82,8 +82,33 @@ IOParameter::GetIOParameterByJson(const JsonType& json) { } catch (std::invalid_argument& error) { return nullptr; } + if (io_ptr != nullptr) { + io_ptr->LoadReadCacheConfig(json); + } return io_ptr; } + +void +IOParameter::LoadReadCacheConfig(const JsonType& json) { + if (json.Contains(READ_CACHE_ENABLED_KEY)) { + CHECK_ARGUMENT(json[READ_CACHE_ENABLED_KEY].IsBool(), + "enable_read_cache must be a boolean"); + enable_read_cache_ = json[READ_CACHE_ENABLED_KEY].GetBool(); + } + if (json.Contains(READ_CACHE_TOTAL_CACHE_SIZE_KEY)) { + const auto& val = json[READ_CACHE_TOTAL_CACHE_SIZE_KEY]; + CHECK_ARGUMENT(val.IsNumberUnsigned(), "total_cache_size must be a non-negative integer"); + read_cache_total_size_ = val.GetUint64(); + } +} + +void +IOParameter::AppendReadCacheConfig(JsonType& json) const { + if (enable_read_cache_) { + json[READ_CACHE_ENABLED_KEY].SetBool(true); + json[READ_CACHE_TOTAL_CACHE_SIZE_KEY].SetUint64(read_cache_total_size_); + } +} IOParameter::IOParameter(std::string name) : name_(std::move(name)) { } } // namespace vsag diff --git a/src/io/common/io_parameter.h b/src/io/common/io_parameter.h index c12c488f76..73d4582438 100644 --- a/src/io/common/io_parameter.h +++ b/src/io/common/io_parameter.h @@ -39,6 +39,15 @@ class IOParameter : public Parameter { static IOParamPtr GetIOParameterByJson(const JsonType& json); + void + LoadReadCacheConfig(const JsonType& json); + + void + AppendReadCacheConfig(JsonType& json) const; + + bool enable_read_cache_{false}; + uint64_t read_cache_total_size_{256ULL * 1024 * 1024}; + public: /** * @brief Returns the type name of this IO parameter. diff --git a/src/io/memory_block_io/memory_block_io.cpp b/src/io/memory_block_io/memory_block_io.cpp index ef27017139..adc6222033 100644 --- a/src/io/memory_block_io/memory_block_io.cpp +++ b/src/io/memory_block_io/memory_block_io.cpp @@ -39,6 +39,7 @@ MemoryBlockIO::MemoryBlockIO(const MemoryBlockIOParamPtr& param, MemoryBlockIO::MemoryBlockIO(const IOParamPtr& param, const IndexCommonParam& common_param) : MemoryBlockIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); } MemoryBlockIO::~MemoryBlockIO() { diff --git a/src/io/memory_block_io/memory_block_io_parameter.cpp b/src/io/memory_block_io/memory_block_io_parameter.cpp index ebbb141977..a0b3bcd522 100644 --- a/src/io/memory_block_io/memory_block_io_parameter.cpp +++ b/src/io/memory_block_io/memory_block_io_parameter.cpp @@ -38,6 +38,7 @@ JsonType MemoryBlockIOParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_BLOCK_MEMORY_IO); + AppendReadCacheConfig(json); return json; } diff --git a/src/io/memory_io/memory_io.h b/src/io/memory_io/memory_io.h index de6c16deeb..1528822341 100644 --- a/src/io/memory_io/memory_io.h +++ b/src/io/memory_io/memory_io.h @@ -73,6 +73,7 @@ class MemoryIO : public BasicIO { */ explicit MemoryIO(const IOParamPtr& param, const IndexCommonParam& common_param) : MemoryIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); } /** diff --git a/src/io/memory_io/memory_io_parameter.cpp b/src/io/memory_io/memory_io_parameter.cpp index a55b940258..c125e642e0 100644 --- a/src/io/memory_io/memory_io_parameter.cpp +++ b/src/io/memory_io/memory_io_parameter.cpp @@ -35,6 +35,7 @@ JsonType MemoryIOParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_MEMORY_IO); + AppendReadCacheConfig(json); return json; } } // namespace vsag diff --git a/src/io/mmap_io/mmap_io.cpp b/src/io/mmap_io/mmap_io.cpp index 95bc5b0f93..490eda66ff 100644 --- a/src/io/mmap_io/mmap_io.cpp +++ b/src/io/mmap_io/mmap_io.cpp @@ -128,18 +128,20 @@ MMapIO::MMapIO(std::string filename, Allocator* allocator) std::error_code(saved_errno, std::system_category()).message())); } this->mapped_ptr_ = static_cast(addr); + this->mapped_size_ = mmap_size; } MMapIO::MMapIO(const MMapIOParamPtr& io_param, const IndexCommonParam& common_param) : MMapIO(io_param->path_, common_param.allocator_.get()){}; MMapIO::MMapIO(const IOParamPtr& param, const IndexCommonParam& common_param) - : MMapIO(std::dynamic_pointer_cast(param), common_param){}; + : MMapIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); +}; MMapIO::~MMapIO() { - auto munmap_size = std::max(this->size_, static_cast(DEFAULT_INIT_MMAP_SIZE)); if (this->mapped_ptr_ != nullptr) { - (void)munmap(this->mapped_ptr_, munmap_size); + (void)munmap(this->mapped_ptr_, this->mapped_size_); } close(this->fd_); // remove file @@ -151,10 +153,7 @@ MMapIO::~MMapIO() { void MMapIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { auto new_size = size + offset; - auto old_size = this->size_; - if (old_size == 0) { - old_size = DEFAULT_INIT_MMAP_SIZE; - } + auto old_size = this->mapped_size_; if (new_size > old_size) { auto ret = IOSyscall::FTruncate(this->fd_, new_size); if (ret == -1) { @@ -176,6 +175,7 @@ MMapIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { } this->mapped_ptr_ = static_cast(new_addr); #endif + this->mapped_size_ = new_size; } this->size_ = std::max(this->size_, new_size); memcpy(this->mapped_ptr_ + offset, data, size); @@ -184,10 +184,7 @@ MMapIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { void MMapIO::ResizeImpl(uint64_t size) { auto new_size = size; - auto old_size = this->size_; - if (old_size == 0) { - old_size = DEFAULT_INIT_MMAP_SIZE; - } + auto old_size = this->mapped_size_; if (new_size > old_size) { auto ret = IOSyscall::FTruncate(this->fd_, new_size); if (ret == -1) { @@ -209,6 +206,7 @@ MMapIO::ResizeImpl(uint64_t size) { } this->mapped_ptr_ = static_cast(new_addr); #endif + this->mapped_size_ = new_size; } else if (new_size < old_size) { auto ret = IOSyscall::FTruncate(this->fd_, new_size); if (ret == -1) { @@ -229,6 +227,7 @@ MMapIO::ResizeImpl(uint64_t size) { } this->mapped_ptr_ = static_cast(new_addr); #endif + this->mapped_size_ = new_size; } this->size_ = new_size; } diff --git a/src/io/mmap_io/mmap_io.h b/src/io/mmap_io/mmap_io.h index 814610b957..be47f83e77 100644 --- a/src/io/mmap_io/mmap_io.h +++ b/src/io/mmap_io/mmap_io.h @@ -38,6 +38,9 @@ class MMapIO : public BasicIO { /// Indicates deserialization is required when loading from disk. static constexpr bool SkipDeserialize = false; + // MMapIO keeps a non-empty backing mapping for subsequent writes. + static constexpr bool SupportsZeroSizeResize = false; + public: /** * @brief Constructs a MMapIO object with a filename and allocator. @@ -133,6 +136,9 @@ class MMapIO : public BasicIO { /// Pointer to the base of the memory-mapped region. uint8_t* mapped_ptr_{nullptr}; + /// Physical extent of mapped_ptr_, which can differ from the logical size. + uint64_t mapped_size_{0}; + /// Flag indicating if file existed before opening; false means file will be removed on destruction. bool exist_file_{false}; }; diff --git a/src/io/mmap_io/mmap_io_parameter.cpp b/src/io/mmap_io/mmap_io_parameter.cpp index dd69f7a429..10035d449e 100644 --- a/src/io/mmap_io/mmap_io_parameter.cpp +++ b/src/io/mmap_io/mmap_io_parameter.cpp @@ -37,6 +37,7 @@ MMapIOParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_MMAP_IO); json[IO_FILE_PATH_KEY].SetString(this->path_); + AppendReadCacheConfig(json); return json; } } // namespace vsag diff --git a/src/io/noncontinuous_io/noncontinuous_io.h b/src/io/noncontinuous_io/noncontinuous_io.h index 9775f2f5b4..f0d1855479 100644 --- a/src/io/noncontinuous_io/noncontinuous_io.h +++ b/src/io/noncontinuous_io/noncontinuous_io.h @@ -91,7 +91,7 @@ class NonContinuousIO : public BasicIO> { while (cur_size < size) { auto area = start_area->first; auto area_size = std::min(size - cur_size, area.size - (start_offset - area.offset)); - inner_io_->WriteImpl(data + cur_size, area_size, start_offset); + inner_io_->Write(data + cur_size, area_size, start_offset); cur_size += area_size; start_area++; if (start_area != areas_.end()) { @@ -134,7 +134,7 @@ class NonContinuousIO : public BasicIO> { start_offset = start_area->first.offset; } } - ret = inner_io_->MultiReadImpl(data, sizes.data(), offsets.data(), sizes.size()); + ret = inner_io_->MultiRead(data, sizes.data(), offsets.data(), sizes.size()); return ret; } diff --git a/src/io/read_cache/lru_page_cache.cpp b/src/io/read_cache/lru_page_cache.cpp new file mode 100644 index 0000000000..fd4d1f7814 --- /dev/null +++ b/src/io/read_cache/lru_page_cache.cpp @@ -0,0 +1,53 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/lru_page_cache.h" + +namespace vsag { + +LRUPageCache::LRUPageCache(uint64_t max_pages) : PageCache(max_pages) { +} + +void +LRUPageCache::OnAccess(uint64_t page_id) { + auto it = iters_.find(page_id); + if (it != iters_.end()) { + order_.splice(order_.begin(), order_, it->second); + } +} + +void +LRUPageCache::OnInsert(uint64_t page_id) { + order_.push_front(page_id); + iters_[page_id] = order_.begin(); +} + +void +LRUPageCache::OnRemove(uint64_t page_id) { + auto it = iters_.find(page_id); + if (it != iters_.end()) { + order_.erase(it->second); + iters_.erase(it); + } +} + +uint64_t +LRUPageCache::PickVictim() { + if (order_.empty()) { + return UINT64_MAX; + } + return order_.back(); +} + +} // namespace vsag diff --git a/src/io/read_cache/lru_page_cache.h b/src/io/read_cache/lru_page_cache.h new file mode 100644 index 0000000000..263a7d46bc --- /dev/null +++ b/src/io/read_cache/lru_page_cache.h @@ -0,0 +1,50 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/page_cache.h" + +namespace vsag { + +/** + * @brief PageCache with Least-Recently-Used eviction policy. + */ +class LRUPageCache : public PageCache { +public: + explicit LRUPageCache(uint64_t max_pages); + +protected: + void + OnAccess(uint64_t page_id) override; + + void + OnInsert(uint64_t page_id) override; + + void + OnRemove(uint64_t page_id) override; + + uint64_t + PickVictim() override; + +private: + std::list order_; + std::unordered_map::iterator> iters_; +}; + +} // namespace vsag diff --git a/src/io/read_cache/lru_page_cache_test.cpp b/src/io/read_cache/lru_page_cache_test.cpp new file mode 100644 index 0000000000..9fb2a160b1 --- /dev/null +++ b/src/io/read_cache/lru_page_cache_test.cpp @@ -0,0 +1,72 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/lru_page_cache.h" + +#include "impl/allocator/safe_allocator.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +PagePtr +MakePage(Allocator* allocator, uint8_t value) { + auto page = std::make_shared(allocator); + page->Data()[0] = value; + return page; +} + +} // namespace + +TEST_CASE("LRUPageCache Evicts Least Recently Used Page", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + LRUPageCache cache(2); + cache.Insert(1, MakePage(allocator.get(), 1)); + cache.Insert(2, MakePage(allocator.get(), 2)); + + REQUIRE(cache.Get(1) != nullptr); + cache.Insert(3, MakePage(allocator.get(), 3)); + + REQUIRE(cache.Size() == 2); + REQUIRE(cache.Get(1) != nullptr); + REQUIRE(cache.Get(2) == nullptr); + REQUIRE(cache.Get(3) != nullptr); +} + +TEST_CASE("LRUPageCache Duplicate Insert Keeps Original Page", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + LRUPageCache cache(2); + auto first = MakePage(allocator.get(), 1); + auto second = MakePage(allocator.get(), 2); + + REQUIRE(cache.Insert(1, first) == first); + REQUIRE(cache.Insert(1, second) == first); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(1)->Data()[0] == 1); +} + +TEST_CASE("LRUPageCache Remove Updates Eviction State", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + LRUPageCache cache(2); + cache.Insert(1, MakePage(allocator.get(), 1)); + cache.Insert(2, MakePage(allocator.get(), 2)); + cache.Remove(1); + cache.Insert(3, MakePage(allocator.get(), 3)); + + REQUIRE(cache.Size() == 2); + REQUIRE(cache.Get(1) == nullptr); + REQUIRE(cache.Get(2) != nullptr); + REQUIRE(cache.Get(3) != nullptr); +} diff --git a/src/io/read_cache/page.h b/src/io/read_cache/page.h new file mode 100644 index 0000000000..ba2cb354c8 --- /dev/null +++ b/src/io/read_cache/page.h @@ -0,0 +1,58 @@ +// Copyright 2024-present the vsag project +// +// 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 "vsag/allocator.h" + +namespace vsag { + +class Page { +public: + explicit Page(Allocator* allocator) : allocator_(allocator) { + data_ = static_cast(allocator_->Allocate(Page::DEFAULT_PAGE_SIZE)); + } + + ~Page() { + if (data_ != nullptr) { + allocator_->Deallocate(data_); + } + } + + Page(const Page&) = delete; + Page& + operator=(const Page&) = delete; + + [[nodiscard]] uint8_t* + Data() { + return data_; + } + [[nodiscard]] const uint8_t* + Data() const { + return data_; + } + + static constexpr uint64_t DEFAULT_PAGE_SIZE = 128 * 1024; + +private: + uint8_t* data_{nullptr}; + Allocator* allocator_{nullptr}; +}; + +using PagePtr = std::shared_ptr; + +} // namespace vsag diff --git a/src/io/read_cache/page_cache.cpp b/src/io/read_cache/page_cache.cpp new file mode 100644 index 0000000000..7ae0dc69b8 --- /dev/null +++ b/src/io/read_cache/page_cache.cpp @@ -0,0 +1,84 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/page_cache.h" + +#include + +namespace vsag { + +PageCache::PageCache(uint64_t max_pages) : max_pages_(max_pages) { +} + +PagePtr +PageCache::Get(uint64_t page_id) { + std::scoped_lock lock(mutex_); + auto it = pages_.find(page_id); + if (it == pages_.end()) { + return nullptr; + } + OnAccess(page_id); + return it->second; +} + +PagePtr +PageCache::Insert(uint64_t page_id, PagePtr page) { + std::scoped_lock lock(mutex_); + auto existing = pages_.find(page_id); + if (existing != pages_.end()) { + OnAccess(page_id); + return existing->second; + } + if (max_pages_ == 0) { + return page; + } + while (pages_.size() >= max_pages_) { + uint64_t victim = PickVictim(); + if (victim == UINT64_MAX or pages_.find(victim) == pages_.end()) { + victim = pages_.begin()->first; + } + OnRemove(victim); + pages_.erase(victim); + } + pages_[page_id] = std::move(page); + OnInsert(page_id); + return pages_[page_id]; +} + +void +PageCache::Remove(uint64_t page_id) { + std::scoped_lock lock(mutex_); + auto it = pages_.find(page_id); + if (it != pages_.end()) { + OnRemove(page_id); + pages_.erase(it); + } +} + +void +PageCache::Clear() { + std::scoped_lock lock(mutex_); + for (const auto& page_pair : pages_) { + OnRemove(page_pair.first); + } + pages_.clear(); +} + +uint64_t +PageCache::Size() const { + std::scoped_lock lock(mutex_); + return pages_.size(); +} + +} // namespace vsag diff --git a/src/io/read_cache/page_cache.h b/src/io/read_cache/page_cache.h new file mode 100644 index 0000000000..71fce90315 --- /dev/null +++ b/src/io/read_cache/page_cache.h @@ -0,0 +1,99 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/page.h" + +namespace vsag { + +/** + * @brief Page storage manager with pluggable eviction policy. + * + * PageCache owns the cached pages and decides which page to evict + * when the capacity is reached. Pages are managed by shared_ptr: + * a page returned to the caller stays valid even if it is evicted + * from the cache concurrently. + */ +class PageCache { +public: + explicit PageCache(uint64_t max_pages); + + virtual ~PageCache() = default; + + /** + * @brief Look up a page, marking it as recently accessed. + * + * @param page_id The page to look up. + * @return The cached page, or nullptr on miss. + */ + virtual PagePtr + Get(uint64_t page_id); + + /** + * @brief Insert a page, evicting victims first if the cache is full. + * + * If the page already exists, the existing page is returned instead. + * + * @param page_id The page id. + * @param page The page to insert. + * @return The cached page. + */ + virtual PagePtr + Insert(uint64_t page_id, PagePtr page); + + /** + * @brief Remove a page from the cache. + * + * @param page_id The page to remove. + */ + virtual void + Remove(uint64_t page_id); + + /** + * @brief Clear all cached pages. + */ + virtual void + Clear(); + + /** + * @brief Number of pages currently cached. + */ + virtual uint64_t + Size() const; + +protected: + virtual void + OnAccess(uint64_t page_id) = 0; + + virtual void + OnInsert(uint64_t page_id) = 0; + + virtual void + OnRemove(uint64_t page_id) = 0; + + virtual uint64_t + PickVictim() = 0; + +protected: + mutable std::mutex mutex_; + std::unordered_map pages_; + uint64_t max_pages_{0}; +}; + +} // namespace vsag diff --git a/src/io/read_cache/page_cache_test.cpp b/src/io/read_cache/page_cache_test.cpp new file mode 100644 index 0000000000..dc229b9825 --- /dev/null +++ b/src/io/read_cache/page_cache_test.cpp @@ -0,0 +1,114 @@ +// Copyright 2024-present the vsag project +// +// 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 "io/read_cache/page_cache.h" + +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +class FifoPageCache : public PageCache { +public: + explicit FifoPageCache(uint64_t max_pages) : PageCache(max_pages) { + } + +protected: + void + OnAccess(uint64_t page_id) override { + last_access_ = page_id; + } + + void + OnInsert(uint64_t page_id) override { + order_.push_back(page_id); + } + + void + OnRemove(uint64_t page_id) override { + auto it = std::find(order_.begin(), order_.end(), page_id); + if (it != order_.end()) { + order_.erase(it); + } + } + + uint64_t + PickVictim() override { + if (order_.empty()) { + return UINT64_MAX; + } + return order_.front(); + } + +public: + uint64_t last_access_{UINT64_MAX}; + +private: + std::deque order_; +}; + +PagePtr +MakePage(Allocator* allocator, uint8_t value) { + auto page = std::make_shared(allocator); + page->Data()[0] = value; + return page; +} + +} // namespace + +TEST_CASE("PageCache Insert Get Remove Test", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + FifoPageCache cache(2); + auto page = MakePage(allocator.get(), 7); + + REQUIRE(cache.Insert(1, page) == page); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(1) == page); + REQUIRE(cache.last_access_ == 1); + REQUIRE(cache.Get(2) == nullptr); + + cache.Remove(1); + REQUIRE(cache.Size() == 0); + REQUIRE(cache.Get(1) == nullptr); +} + +TEST_CASE("PageCache Duplicate Insert Test", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + FifoPageCache cache(2); + auto first = MakePage(allocator.get(), 1); + auto second = MakePage(allocator.get(), 2); + + REQUIRE(cache.Insert(10, first) == first); + REQUIRE(cache.Insert(10, second) == first); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(10)->Data()[0] == 1); +} + +TEST_CASE("PageCache Eviction Test", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + FifoPageCache cache(2); + cache.Insert(1, MakePage(allocator.get(), 1)); + cache.Insert(2, MakePage(allocator.get(), 2)); + cache.Insert(3, MakePage(allocator.get(), 3)); + + REQUIRE(cache.Size() == 2); + REQUIRE(cache.Get(1) == nullptr); + REQUIRE(cache.Get(2) != nullptr); + REQUIRE(cache.Get(3) != nullptr); +} diff --git a/src/io/read_cache/read_cache_test.cpp b/src/io/read_cache/read_cache_test.cpp new file mode 100644 index 0000000000..36e5faea68 --- /dev/null +++ b/src/io/read_cache/read_cache_test.cpp @@ -0,0 +1,155 @@ +// Copyright 2024-present the vsag project +// +// 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 +#include + +#include "impl/allocator/safe_allocator.h" +#include "io/buffer_io/buffer_io.h" +#include "io/buffer_io/buffer_io_parameter.h" +#include "io/common/basic_io_test.h" +#include "io/common/io_parameter.h" +#include "io/memory_io/memory_io_parameter.h" +#include "io/read_cache/page.h" +#include "io/reader_io/reader_io.h" +#include "io/reader_io/reader_io_parameter.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +IOParamPtr +MakeReadCacheParam(uint64_t page_count = 4) { + auto param = std::make_shared(); + param->enable_read_cache_ = true; + param->read_cache_total_size_ = Page::DEFAULT_PAGE_SIZE * page_count; + return param; +} + +class TestReader : public Reader { +public: + explicit TestReader(const std::vector& data) : data_(data) { + } + + void + Read(uint64_t offset, uint64_t len, void* dest) override { + std::memcpy(dest, data_.data() + offset, len); + } + + void + AsyncRead(uint64_t offset, uint64_t len, void* dest, CallBack callback) override { + Read(offset, len, dest); + callback(IOErrorCode::IO_SUCCESS, "success"); + } + + uint64_t + Size() const override { + return data_.size(); + } + +private: + const std::vector& data_; +}; + +} // namespace + +TEST_CASE("BasicIO cache component basic test", "[ReadCache][ut]") { + fixtures::TempDir dir("read_cache_basic"); + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + BufferIO io(dir.GenerateRandomFile(false), allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + TestBasicReadWrite(io); +} + +TEST_CASE("BasicIO cache component file backend test", "[ReadCache][ut]") { + fixtures::TempDir dir("read_cache_buffer"); + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + BufferIO io(dir.GenerateRandomFile(false), allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + TestBasicReadWrite(io); +} + +TEST_CASE("BasicIO cache component invalidates writes", "[ReadCache][ut]") { + fixtures::TempDir dir("read_cache_invalidate"); + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + BufferIO io(dir.GenerateRandomFile(false), allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector old_data(128, 0x11); + std::vector new_data(64, 0x22); + std::vector read_buf(64); + io.Write(old_data.data(), old_data.size(), 0); + REQUIRE(io.Read(read_buf.size(), 32, read_buf.data())); + + io.Write(new_data.data(), new_data.size(), 32); + REQUIRE(io.Read(read_buf.size(), 32, read_buf.data())); + REQUIRE(std::memcmp(read_buf.data(), new_data.data(), read_buf.size()) == 0); +} + +TEST_CASE("BasicIO cache component direct and multi read", "[ReadCache][ut]") { + fixtures::TempDir dir("read_cache_multi_read"); + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + BufferIO io(dir.GenerateRandomFile(false), allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + uint64_t offset = Page::DEFAULT_PAGE_SIZE - 30; + std::vector data(100, 0xCD); + io.Write(data.data(), data.size(), offset); + + bool need_release = false; + const auto* ptr = io.Read(data.size(), offset, need_release); + REQUIRE(ptr != nullptr); + REQUIRE(need_release); + REQUIRE(std::memcmp(ptr, data.data(), data.size()) == 0); + io.Release(ptr); + + uint64_t sizes[] = {30, 40}; + uint64_t offsets[] = {offset, offset + 30}; + std::vector result(70); + REQUIRE(io.MultiRead(result.data(), sizes, offsets, 2)); + REQUIRE(std::memcmp(result.data(), data.data(), result.size()) == 0); +} + +TEST_CASE("BasicIO cache component initializes ReaderIO", "[ReadCache][ut]") { + std::vector data(1024); + for (uint64_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i); + } + + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + ReaderIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + auto reader_param = std::make_shared(); + reader_param->reader = std::make_shared(data); + io.InitIO(reader_param); + + std::vector read_buf(data.size()); + REQUIRE(io.Read(read_buf.size(), 0, read_buf.data())); + REQUIRE(read_buf == data); +} + +TEST_CASE("ReadCache configuration test", "[ReadCache][ut]") { + JsonType json; + json["type"].SetString("memory_io"); + json["enable_read_cache"].SetBool(true); + json["total_cache_size"].SetUint64(4096); + + auto param = IOParameter::GetIOParameterByJson(json); + REQUIRE(param != nullptr); + REQUIRE(param->enable_read_cache_); + REQUIRE(param->read_cache_total_size_ == 4096); + REQUIRE(param->ToJson()["enable_read_cache"].GetBool()); + REQUIRE(param->ToJson()["total_cache_size"].GetUint64() == 4096); +} diff --git a/src/io/reader_io/reader_io.cpp b/src/io/reader_io/reader_io.cpp index 7be91d6630..fac8ccd33f 100644 --- a/src/io/reader_io/reader_io.cpp +++ b/src/io/reader_io/reader_io.cpp @@ -31,6 +31,7 @@ ReaderIO::ReaderIO(const ReaderIOParamPtr& /*param*/, const IndexCommonParam& co ReaderIO::ReaderIO(const IOParamPtr& param, const IndexCommonParam& common_param) : ReaderIO(std::dynamic_pointer_cast(param), common_param) { + EnableReadCache(param); } void @@ -45,7 +46,13 @@ ReaderIO::InitIOImpl(const vsag::IOParamPtr& io_param) { throw VsagException(ErrorType::INTERNAL_ERROR, "ReaderIOParam is required for ReaderIO initialization."); } + if (reader_param->reader == nullptr) { + throw VsagException(ErrorType::INTERNAL_ERROR, "ReaderIO requires a non-null reader."); + } reader_ = reader_param->reader; + if (not HasDeserialized()) { + this->size_ = reader_->Size(); + } } bool diff --git a/src/io/reader_io/reader_io_parameter.h b/src/io/reader_io/reader_io_parameter.h index f0f6210b63..0f6a92348e 100644 --- a/src/io/reader_io/reader_io_parameter.h +++ b/src/io/reader_io/reader_io_parameter.h @@ -34,6 +34,7 @@ class ReaderIOParameter : public IOParameter { ToJson() const override { JsonType json; json[TYPE_KEY].SetString(IO_TYPE_VALUE_READER_IO); + AppendReadCacheConfig(json); return json; } diff --git a/src/io/uring_io/uring_io.cpp b/src/io/uring_io/uring_io.cpp index 83a0a8dddd..ef844e6b0d 100644 --- a/src/io/uring_io/uring_io.cpp +++ b/src/io/uring_io/uring_io.cpp @@ -81,6 +81,7 @@ UringIO::UringIO(const IOParamPtr& param, const IndexCommonParam& common_param) return p; }(), common_param) { + EnableReadCache(param); } UringIO::~UringIO() { diff --git a/src/io/uring_io/uring_io_parameter.cpp b/src/io/uring_io/uring_io_parameter.cpp index 967026e0a2..c7bfbe3b06 100644 --- a/src/io/uring_io/uring_io_parameter.cpp +++ b/src/io/uring_io/uring_io_parameter.cpp @@ -49,6 +49,7 @@ UringIOParameter::ToJson() const { json[TYPE_KEY].SetString(IO_TYPE_VALUE_URING_IO); json[IO_FILE_PATH_KEY].SetString(this->path_); json[IO_DIRECT_READ_KEY].SetBool(this->direct_read_); + AppendReadCacheConfig(json); return json; } } // namespace vsag diff --git a/tests/test_hgraph.cpp b/tests/test_hgraph.cpp index 8d8c326342..bf07cc2fff 100644 --- a/tests/test_hgraph.cpp +++ b/tests/test_hgraph.cpp @@ -3858,3 +3858,43 @@ TEST_CASE("HGraph Concurrent Tune(disable_future_tuning=false) and CalDistanceBy REQUIRE(cal_count.load() > 0); } + +TEST_CASE_PERSISTENT_FIXTURE(fixtures::HGraphTestIndex, + "HGraph ReadCache With Async Backend", + "[ft][hgraph][cacheio][pr]") { + auto resource = HGraphTestIndex::GetResource(true); + auto dim = 128; + std::string search_param = R"({"hgraph": {"ef_search": 200}})"; + + auto dataset = HGraphTestIndex::pool.GetDatasetAndCreate(dim, resource->base_count, "l2"); + + auto precise_file_path = HGraphTestIndex::dir.GenerateRandomFile(false); + std::string read_cache_param = fmt::format(R"( + {{ + "dtype": "float32", + "metric_type": "l2", + "dim": {}, + "index_param": {{ + "use_reorder": true, + "base_quantization_type": "sq8", + "precise_quantization_type": "fp32", + "max_degree": 96, + "ef_construction": 500, + "base_pq_dim": {}, + "base_io_type": "memory_io", + "precise_io_type": "async_io", + "precise_enable_read_cache": true, + "precise_file_path": "{}", + "precise_cache_total_size": 268435456, + "graph_io_type": "block_memory_io" + }} + }} + )", + dim, + dim, + precise_file_path); + + auto cache_index = TestIndex::TestFactory(name, read_cache_param, true); + TestIndex::TestBuildIndex(cache_index, dataset, true); + HGraphTestIndex::TestGeneral(cache_index, dataset, search_param, 0.98f); +} diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 5e71284a72..f6bc867aca 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -591,6 +591,49 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, } } +TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, + "IVF ReadCache With Async Backend", + "[ft][ivf][cacheio][pr]") { + auto dim = 64; + auto precise_file_path = fixtures::IVFTestIndex::dir.GenerateRandomFile(false); + std::string param = fmt::format(R"( + {{ + "dtype": "float32", + "metric_type": "l2", + "dim": {}, + "index_param": {{ + "buckets_count": 32, + "base_quantization_type": "fp32", + "partition_strategy_type": "ivf", + "ivf_train_type": "random", + "train_sample_count": 512, + "use_reorder": true, + "precise_quantization_type": "fp32", + "base_io_type": "memory_io", + "precise_io_type": "async_io", + "precise_enable_read_cache": true, + "precise_file_path": "{}", + "precise_cache_total_size": 131072 + }} + }} + )", + dim, + precise_file_path); + + auto index = fixtures::TestIndex::TestFactory(fixtures::IVFTestIndex::name, param, true); + auto dataset = fixtures::IVFTestIndex::pool.GetDatasetAndCreate(dim, 512, "l2"); + fixtures::TestIndex::TestBuildIndex(index, dataset, true); + + auto serialize_result = index->Serialize(); + REQUIRE(serialize_result.has_value()); + + auto restored = fixtures::TestIndex::TestFactory(fixtures::IVFTestIndex::name, param, true); + REQUIRE(restored->Deserialize(serialize_result.value()).has_value()); + + auto search_param = fmt::format(fixtures::search_param_tmp, 32); + fixtures::IVFTestIndex::TestGeneral(restored, dataset, search_param, 0.90F); +} + TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, "IVF RabitQ base quantization", "[ft][ivf][rabitq]") {