enhance: Make GeometryCache an optional configuration (#45192)

issue: #45187

Signed-off-by: Cai Zhang <cai.zhang@zilliz.com>
This commit is contained in:
cai.zhang 2025-11-03 19:59:32 +08:00 committed by GitHub
parent ae03dee116
commit ed8ba4a28c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 182 additions and 92 deletions

View File

@ -465,6 +465,7 @@ queryNode:
memExpansionRate: 1.15 # extra memory needed by building interim index memExpansionRate: 1.15 # extra memory needed by building interim index
buildParallelRate: 0.5 # the ratio of building interim index parallel matched with cpu num buildParallelRate: 0.5 # the ratio of building interim index parallel matched with cpu num
multipleChunkedEnable: true # Deprecated. Enable multiple chunked search multipleChunkedEnable: true # Deprecated. Enable multiple chunked search
enableGeometryCache: false # Enable geometry cache for geometry data
tieredStorage: tieredStorage:
warmup: warmup:
# options: sync, disable. # options: sync, disable.

View File

@ -114,7 +114,7 @@ class SimpleGeometryCacheManager {
SimpleGeometryCacheManager() = default; SimpleGeometryCacheManager() = default;
SimpleGeometryCache& SimpleGeometryCache&
GetCache(int64_t segment_id, FieldId field_id) { GetOrCreateCache(int64_t segment_id, FieldId field_id) {
std::lock_guard<std::mutex> lock(mutex_); std::lock_guard<std::mutex> lock(mutex_);
auto key = MakeCacheKey(segment_id, field_id); auto key = MakeCacheKey(segment_id, field_id);
auto it = caches_.find(key); auto it = caches_.find(key);
@ -128,6 +128,17 @@ class SimpleGeometryCacheManager {
return *cache_ptr; return *cache_ptr;
} }
SimpleGeometryCache*
GetCache(int64_t segment_id, FieldId field_id) {
std::lock_guard<std::mutex> lock(mutex_);
auto key = MakeCacheKey(segment_id, field_id);
auto it = caches_.find(key);
if (it != caches_.end()) {
return it->second.get();
}
return nullptr;
}
void void
RemoveCache(GEOSContextHandle_t ctx, int64_t segment_id, FieldId field_id) { RemoveCache(GEOSContextHandle_t ctx, int64_t segment_id, FieldId field_id) {
std::lock_guard<std::mutex> lock(mutex_); std::lock_guard<std::mutex> lock(mutex_);
@ -184,26 +195,4 @@ class SimpleGeometryCacheManager {
} // namespace exec } // namespace exec
// Convenient global functions for direct access to geometry cache
inline const Geometry*
GetGeometryByOffset(int64_t segment_id, FieldId field_id, size_t offset) {
auto& cache = exec::SimpleGeometryCacheManager::Instance().GetCache(
segment_id, field_id);
return cache.GetByOffset(offset);
}
inline void
RemoveGeometryCache(GEOSContextHandle_t ctx,
int64_t segment_id,
FieldId field_id) {
exec::SimpleGeometryCacheManager::Instance().RemoveCache(
ctx, segment_id, field_id);
}
inline void
RemoveSegmentGeometryCaches(GEOSContextHandle_t ctx, int64_t segment_id) {
exec::SimpleGeometryCacheManager::Instance().RemoveSegmentCaches(
ctx, segment_id);
}
} // namespace milvus } // namespace milvus

View File

@ -31,11 +31,11 @@ namespace exec {
const Geometry& right_source) { \ const Geometry& right_source) { \
AssertInfo(segment_offsets != nullptr, \ AssertInfo(segment_offsets != nullptr, \
"segment_offsets should not be nullptr"); \ "segment_offsets should not be nullptr"); \
/* Unified path using simple WKB-content-based cache for both sealed and growing segments. */ \ auto* geometry_cache = \
auto& geometry_cache = \
SimpleGeometryCacheManager::Instance().GetCache( \ SimpleGeometryCacheManager::Instance().GetCache( \
this->segment_->get_segment_id(), field_id_); \ this->segment_->get_segment_id(), field_id_); \
auto cache_lock = geometry_cache.AcquireReadLock(); \ if (geometry_cache) { \
auto cache_lock = geometry_cache->AcquireReadLock(); \
for (int i = 0; i < size; ++i) { \ for (int i = 0; i < size; ++i) { \
if (valid_data != nullptr && !valid_data[i]) { \ if (valid_data != nullptr && !valid_data[i]) { \
res[i] = valid_res[i] = false; \ res[i] = valid_res[i] = false; \
@ -43,11 +43,23 @@ namespace exec {
} \ } \
auto absolute_offset = segment_offsets[i]; \ auto absolute_offset = segment_offsets[i]; \
auto cached_geometry = \ auto cached_geometry = \
geometry_cache.GetByOffsetUnsafe(absolute_offset); \ geometry_cache->GetByOffsetUnsafe(absolute_offset); \
AssertInfo(cached_geometry != nullptr, \ AssertInfo(cached_geometry != nullptr, \
"cached geometry is nullptr"); \ "cached geometry is nullptr"); \
res[i] = cached_geometry->method(right_source); \ res[i] = cached_geometry->method(right_source); \
} \ } \
} else { \
GEOSContextHandle_t ctx_ = GEOS_init_r(); \
for (int i = 0; i < size; ++i) { \
if (valid_data != nullptr && !valid_data[i]) { \
res[i] = valid_res[i] = false; \
continue; \
} \
res[i] = Geometry(ctx_, data[i].data(), data[i].size()) \
.method(right_source); \
} \
GEOS_finish_r(ctx_); \
} \
}; \ }; \
int64_t processed_size = ProcessDataChunks<_DataType, true>( \ int64_t processed_size = ProcessDataChunks<_DataType, true>( \
execute_sub_batch, std::nullptr_t{}, res, valid_res, right_source); \ execute_sub_batch, std::nullptr_t{}, res, valid_res, right_source); \
@ -57,7 +69,6 @@ namespace exec {
processed_size, \ processed_size, \
real_batch_size); \ real_batch_size); \
return res_vec; return res_vec;
// Specialized macro for distance-based operations (ST_DWITHIN) // Specialized macro for distance-based operations (ST_DWITHIN)
#define GEOMETRY_EXECUTE_SUB_BATCH_WITH_COMPARISON_DISTANCE(_DataType, method) \ #define GEOMETRY_EXECUTE_SUB_BATCH_WITH_COMPARISON_DISTANCE(_DataType, method) \
auto execute_sub_batch = [this](const _DataType* data, \ auto execute_sub_batch = [this](const _DataType* data, \
@ -70,10 +81,11 @@ namespace exec {
const Geometry& right_source) { \ const Geometry& right_source) { \
AssertInfo(segment_offsets != nullptr, \ AssertInfo(segment_offsets != nullptr, \
"segment_offsets should not be nullptr"); \ "segment_offsets should not be nullptr"); \
auto& geometry_cache = \ auto* geometry_cache = \
SimpleGeometryCacheManager::Instance().GetCache( \ SimpleGeometryCacheManager::Instance().GetCache( \
this->segment_->get_segment_id(), field_id_); \ this->segment_->get_segment_id(), field_id_); \
auto cache_lock = geometry_cache.AcquireReadLock(); \ if (geometry_cache) { \
auto cache_lock = geometry_cache->AcquireReadLock(); \
for (int i = 0; i < size; ++i) { \ for (int i = 0; i < size; ++i) { \
if (valid_data != nullptr && !valid_data[i]) { \ if (valid_data != nullptr && !valid_data[i]) { \
res[i] = valid_res[i] = false; \ res[i] = valid_res[i] = false; \
@ -81,10 +93,23 @@ namespace exec {
} \ } \
auto absolute_offset = segment_offsets[i]; \ auto absolute_offset = segment_offsets[i]; \
auto cached_geometry = \ auto cached_geometry = \
geometry_cache.GetByOffsetUnsafe(absolute_offset); \ geometry_cache->GetByOffsetUnsafe(absolute_offset); \
AssertInfo(cached_geometry != nullptr, \ AssertInfo(cached_geometry != nullptr, \
"cached geometry is nullptr"); \ "cached geometry is nullptr"); \
res[i] = cached_geometry->method(right_source, expr_->distance_); \ res[i] = \
cached_geometry->method(right_source, expr_->distance_); \
} \
} else { \
GEOSContextHandle_t ctx_ = GEOS_init_r(); \
for (int i = 0; i < size; ++i) { \
if (valid_data != nullptr && !valid_data[i]) { \
res[i] = valid_res[i] = false; \
continue; \
} \
res[i] = Geometry(ctx_, data[i].data(), data[i].size()) \
.method(right_source, expr_->distance_); \
} \
GEOS_finish_r(ctx_); \
} \ } \
}; \ }; \
int64_t processed_size = ProcessDataChunks<_DataType, true>( \ int64_t processed_size = ProcessDataChunks<_DataType, true>( \
@ -357,21 +382,22 @@ PhyGISFunctionFilterExpr::EvalForIndexSegment() {
}; };
// Lambda: Process sealed segment data using bulk_subscript with SimpleGeometryCache // Lambda: Process sealed segment data using bulk_subscript with SimpleGeometryCache
auto process_sealed_data = auto process_sealed_data = [&](const std::vector<int64_t>&
[&](const std::vector<int64_t>& hit_offsets) { hit_offsets) {
if (hit_offsets.empty()) if (hit_offsets.empty())
return; return;
// Get simple geometry cache for this segment+field // Get simple geometry cache for this segment+field
auto& geometry_cache = auto* geometry_cache =
SimpleGeometryCacheManager::Instance().GetCache( SimpleGeometryCacheManager::Instance().GetCache(
segment_->get_segment_id(), field_id_); segment_->get_segment_id(), field_id_);
auto cache_lock = geometry_cache.AcquireReadLock(); if (geometry_cache) {
auto cache_lock = geometry_cache->AcquireReadLock();
for (size_t i = 0; i < hit_offsets.size(); ++i) { for (size_t i = 0; i < hit_offsets.size(); ++i) {
const auto pos = hit_offsets[i]; const auto pos = hit_offsets[i];
auto cached_geometry = auto cached_geometry =
geometry_cache.GetByOffsetUnsafe(pos); geometry_cache->GetByOffsetUnsafe(pos);
// skip invalid geometry // skip invalid geometry
if (cached_geometry == nullptr) { if (cached_geometry == nullptr) {
continue; continue;
@ -383,6 +409,35 @@ PhyGISFunctionFilterExpr::EvalForIndexSegment() {
refined.set(pos); refined.set(pos);
} }
} }
} else {
milvus::OpContext op_ctx;
auto data_array = segment_->bulk_subscript(
&op_ctx, field_id_, hit_offsets.data(), hit_offsets.size());
auto geometry_array =
static_cast<const milvus::proto::schema::GeometryArray*>(
&data_array->scalars().geometry_data());
const auto& valid_data = data_array->valid_data();
GEOSContextHandle_t ctx = GEOS_init_r();
for (size_t i = 0; i < hit_offsets.size(); ++i) {
const auto pos = hit_offsets[i];
// Skip invalid data
if (!valid_data.empty() && !valid_data[i]) {
continue;
}
const auto& wkb_data = geometry_array->data(i);
Geometry left(ctx, wkb_data.data(), wkb_data.size());
bool result = evaluate_geometry(left, query_geometry);
if (result) {
refined.set(pos);
}
}
GEOS_finish_r(ctx);
}
}; };
auto hit_offsets = collect_hits(); auto hit_offsets = collect_hits();

View File

@ -2654,7 +2654,8 @@ ChunkedSegmentSealedImpl::load_field_data_common(
column->ManualEvictCache(); column->ManualEvictCache();
} }
} }
if (data_type == DataType::GEOMETRY) { if (data_type == DataType::GEOMETRY &&
segcore_config_.get_enable_geometry_cache()) {
// Construct GeometryCache for the entire field // Construct GeometryCache for the entire field
LoadGeometryCache(field_id, column); LoadGeometryCache(field_id, column);
} }
@ -2787,8 +2788,8 @@ ChunkedSegmentSealedImpl::LoadGeometryCache(
try { try {
// Get geometry cache for this segment+field // Get geometry cache for this segment+field
auto& geometry_cache = auto& geometry_cache =
milvus::exec::SimpleGeometryCacheManager::Instance().GetCache( milvus::exec::SimpleGeometryCacheManager::Instance()
get_segment_id(), field_id); .GetOrCreateCache(get_segment_id(), field_id);
// Iterate through all chunks and collect WKB data // Iterate through all chunks and collect WKB data
auto num_chunks = column->num_chunks(); auto num_chunks = column->num_chunks();

View File

@ -147,6 +147,16 @@ class SegcoreConfig {
return refine_with_quant_flag_; return refine_with_quant_flag_;
} }
void
set_enable_geometry_cache(bool enable_geometry_cache) {
enable_geometry_cache_ = enable_geometry_cache;
}
bool
get_enable_geometry_cache() const {
return enable_geometry_cache_;
}
private: private:
inline static const std::unordered_set<std::string> inline static const std::unordered_set<std::string>
valid_dense_vector_index_type = { valid_dense_vector_index_type = {
@ -165,6 +175,7 @@ class SegcoreConfig {
inline static knowhere::RefineType refine_type_ = inline static knowhere::RefineType refine_type_ =
knowhere::RefineType::DATA_VIEW; knowhere::RefineType::DATA_VIEW;
inline static bool refine_with_quant_flag_ = false; inline static bool refine_with_quant_flag_ = false;
inline static bool enable_geometry_cache_ = false;
}; };
} // namespace milvus::segcore } // namespace milvus::segcore

View File

@ -226,7 +226,8 @@ SegmentGrowingImpl::Insert(int64_t reserved_offset,
} }
// Build geometry cache for GEOMETRY fields // Build geometry cache for GEOMETRY fields
if (field_meta.get_data_type() == DataType::GEOMETRY) { if (field_meta.get_data_type() == DataType::GEOMETRY &&
segcore_config_.get_enable_geometry_cache()) {
BuildGeometryCacheForInsert( BuildGeometryCacheForInsert(
field_id, field_id,
&insert_record_proto->fields_data(data_offset), &insert_record_proto->fields_data(data_offset),
@ -527,7 +528,8 @@ SegmentGrowingImpl::load_column_group_data_internal(
num_rows); num_rows);
// Build geometry cache for GEOMETRY fields // Build geometry cache for GEOMETRY fields
if (schema_->operator[](field_id).get_data_type() == if (schema_->operator[](field_id).get_data_type() ==
DataType::GEOMETRY) { DataType::GEOMETRY &&
segcore_config_.get_enable_geometry_cache()) {
BuildGeometryCacheForLoad(field_id, field_data); BuildGeometryCacheForLoad(field_id, field_data);
} }
} }
@ -1379,8 +1381,8 @@ SegmentGrowingImpl::BuildGeometryCacheForInsert(FieldId field_id,
try { try {
// Get geometry cache for this segment+field // Get geometry cache for this segment+field
auto& geometry_cache = auto& geometry_cache =
milvus::exec::SimpleGeometryCacheManager::Instance().GetCache( milvus::exec::SimpleGeometryCacheManager::Instance()
get_segment_id(), field_id); .GetOrCreateCache(get_segment_id(), field_id);
// Process geometry data from DataArray // Process geometry data from DataArray
const auto& geometry_data = data_array->scalars().geometry_data(); const auto& geometry_data = data_array->scalars().geometry_data();
@ -1423,8 +1425,8 @@ SegmentGrowingImpl::BuildGeometryCacheForLoad(
try { try {
// Get geometry cache for this segment+field // Get geometry cache for this segment+field
auto& geometry_cache = auto& geometry_cache =
milvus::exec::SimpleGeometryCacheManager::Instance().GetCache( milvus::exec::SimpleGeometryCacheManager::Instance()
get_segment_id(), field_id); .GetOrCreateCache(get_segment_id(), field_id);
// Process each field data chunk // Process each field data chunk
for (const auto& data : field_data) { for (const auto& data : field_data) {

View File

@ -41,6 +41,13 @@ SegcoreSetEnableInterminSegmentIndex(const bool value) {
config.set_enable_interim_segment_index(value); config.set_enable_interim_segment_index(value);
} }
extern "C" void
SegcoreSetEnableGeometryCache(const bool value) {
milvus::segcore::SegcoreConfig& config =
milvus::segcore::SegcoreConfig::default_config();
config.set_enable_geometry_cache(value);
}
extern "C" void extern "C" void
SegcoreSetNlist(const int64_t value) { SegcoreSetNlist(const int64_t value) {
milvus::segcore::SegcoreConfig& config = milvus::segcore::SegcoreConfig& config =

View File

@ -30,6 +30,9 @@ SegcoreSetChunkRows(const int64_t);
void void
SegcoreSetEnableInterminSegmentIndex(const bool); SegcoreSetEnableInterminSegmentIndex(const bool);
void
SegcoreSetEnableGeometryCache(const bool);
void void
SegcoreSetNlist(const int64_t); SegcoreSetNlist(const int64_t);

View File

@ -562,6 +562,12 @@ func InitInterminIndexConfig(params *paramtable.ComponentParam) error {
return HandleCStatus(&status, "InitInterminIndexConfig failed") return HandleCStatus(&status, "InitInterminIndexConfig failed")
} }
func InitGeometryCache(params *paramtable.ComponentParam) error {
enableGeometryCache := C.bool(params.QueryNodeCfg.EnableGeometryCache.GetAsBool())
C.SegcoreSetEnableGeometryCache(enableGeometryCache)
return nil
}
func CleanRemoteChunkManager() { func CleanRemoteChunkManager() {
C.CleanRemoteChunkManagerSingleton() C.CleanRemoteChunkManagerSingleton()
} }

View File

@ -174,6 +174,11 @@ func doInitQueryNodeOnce(ctx context.Context) error {
return err return err
} }
err = InitGeometryCache(paramtable.Get())
if err != nil {
return err
}
InitTraceConfig(paramtable.Get()) InitTraceConfig(paramtable.Get())
C.InitExecExpressionFunctionFactory() C.InitExecExpressionFunctionFactory()

View File

@ -3052,6 +3052,7 @@ type queryNodeConfig struct {
InterimIndexMemExpandRate ParamItem `refreshable:"false"` InterimIndexMemExpandRate ParamItem `refreshable:"false"`
InterimIndexBuildParallelRate ParamItem `refreshable:"false"` InterimIndexBuildParallelRate ParamItem `refreshable:"false"`
MultipleChunkedEnable ParamItem `refreshable:"false"` // Deprecated MultipleChunkedEnable ParamItem `refreshable:"false"` // Deprecated
EnableGeometryCache ParamItem `refreshable:"false"`
// TODO(tiered storage 2) this should be refreshable? // TODO(tiered storage 2) this should be refreshable?
TieredWarmupScalarField ParamItem `refreshable:"false"` TieredWarmupScalarField ParamItem `refreshable:"false"`
@ -3627,6 +3628,15 @@ This defaults to true, indicating that Milvus creates temporary index for growin
} }
p.MultipleChunkedEnable.Init(base.mgr) p.MultipleChunkedEnable.Init(base.mgr)
p.EnableGeometryCache = ParamItem{
Key: "queryNode.segcore.enableGeometryCache",
Version: "2.6.5",
DefaultValue: "false",
Doc: "Enable geometry cache for geometry data",
Export: true,
}
p.EnableGeometryCache.Init(base.mgr)
p.InterimIndexNProbe = ParamItem{ p.InterimIndexNProbe = ParamItem{
Key: "queryNode.segcore.interimIndex.nprobe", Key: "queryNode.segcore.interimIndex.nprobe",
Version: "2.0.0", Version: "2.0.0",