diff --git a/client/index/common.go b/client/index/common.go index 214abdb8ce8..654c42239da 100644 --- a/client/index/common.go +++ b/client/index/common.go @@ -65,4 +65,5 @@ const ( Sorted IndexType = "STL_SORT" Inverted IndexType = "INVERTED" BITMAP IndexType = "BITMAP" + RTREE IndexType = "RTREE" ) diff --git a/client/index/rtree.go b/client/index/rtree.go new file mode 100644 index 00000000000..a01c7c4995a --- /dev/null +++ b/client/index/rtree.go @@ -0,0 +1,70 @@ +// 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. + +package index + +var _ Index = rtreeIndex{} + +// rtreeIndex represents an RTree index for geometry fields +type rtreeIndex struct { + baseIndex +} + +func (idx rtreeIndex) Params() map[string]string { + params := map[string]string{ + IndexTypeKey: string(RTREE), + } + return params +} + +// NewRTreeIndex creates a new RTree index with default parameters +func NewRTreeIndex() Index { + return rtreeIndex{ + baseIndex: baseIndex{ + indexType: RTREE, + }, + } +} + +// NewRTreeIndexWithParams creates a new RTree index with custom parameters +func NewRTreeIndexWithParams() Index { + return rtreeIndex{ + baseIndex: baseIndex{ + indexType: RTREE, + }, + } +} + +// RTreeIndexBuilder provides a fluent API for building RTree indexes +type RTreeIndexBuilder struct { + index rtreeIndex +} + +// NewRTreeIndexBuilder creates a new RTree index builder +func NewRTreeIndexBuilder() *RTreeIndexBuilder { + return &RTreeIndexBuilder{ + index: rtreeIndex{ + baseIndex: baseIndex{ + indexType: RTREE, + }, + }, + } +} + +// Build returns the constructed RTree index +func (b *RTreeIndexBuilder) Build() Index { + return b.index +} diff --git a/client/index/rtree_test.go b/client/index/rtree_test.go new file mode 100644 index 00000000000..7f0bcfc3881 --- /dev/null +++ b/client/index/rtree_test.go @@ -0,0 +1,77 @@ +// 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. + +package index + +import ( + "testing" + + "github.com/stretchr/testify/suite" +) + +type RTreeIndexSuite struct { + suite.Suite +} + +func (s *RTreeIndexSuite) TestNewRTreeIndex() { + idx := NewRTreeIndex() + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) +} + +func (s *RTreeIndexSuite) TestNewRTreeIndexWithParams() { + idx := NewRTreeIndexWithParams() + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilder() { + idx := NewRTreeIndexBuilder(). + Build() + + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilderDefaults() { + idx := NewRTreeIndexBuilder().Build() + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilderChaining() { + builder := NewRTreeIndexBuilder() + + // Test method chaining + result := builder.Build() + + s.Equal(RTREE, result.IndexType()) + + params := result.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) +} + +func TestRTreeIndex(t *testing.T) { + suite.Run(t, new(RTreeIndexSuite)) +} diff --git a/internal/core/CMakeLists.txt b/internal/core/CMakeLists.txt index a7a835f4627..a695151b99f 100644 --- a/internal/core/CMakeLists.txt +++ b/internal/core/CMakeLists.txt @@ -268,6 +268,8 @@ if ( BUILD_DISK_ANN STREQUAL "ON" ) ADD_DEFINITIONS(-DBUILD_DISK_ANN=${BUILD_DISK_ANN}) endif () +ADD_DEFINITIONS(-DBOOST_GEOMETRY_INDEX_DETAIL_EXPERIMENTAL) + # Warning: add_subdirectory(src) must be after append_flags("-ftest-coverage"), # otherwise cpp code coverage tool will miss src folder add_subdirectory( thirdparty ) diff --git a/internal/core/conanfile.py b/internal/core/conanfile.py index 5cb7b9d1b7c..d123b5e26e3 100644 --- a/internal/core/conanfile.py +++ b/internal/core/conanfile.py @@ -6,7 +6,7 @@ class MilvusConan(ConanFile): settings = "os", "compiler", "build_type", "arch" requires = ( "rocksdb/6.29.5@milvus/dev#b1842a53ddff60240c5282a3da498ba1", - "boost/1.82.0#744a17160ebb5838e9115eab4d6d0c06", + "boost/1.83.0@", "onetbb/2021.9.0#4a223ff1b4025d02f31b65aedf5e7f4a", "nlohmann_json/3.11.3#ffb9e9236619f1c883e36662f944345d", "zstd/1.5.5#34e9debe03bf0964834a09dfbc31a5dd", @@ -53,8 +53,8 @@ class MilvusConan(ConanFile): "proj/9.3.1#38e8bacd0f98467d38e20f46a085b4b3", "libtiff/4.6.0#32ca1d04c9f024637d49c0c2882cfdbe", "libgeotiff/1.7.1#0375633ef1116fc067b3773be7fd902f", - "geos/3.12.0#b76c27884c1fa4ee8c9e486337b7dc4e", - "gdal/3.5.3#61a42c933d3440a449cac89fd0866621" + "geos/3.12.0#0b177c90c25a8ca210578fb9e2899c37", + "gdal/3.5.3#61a42c933d3440a449cac89fd0866621", ) generators = ("cmake", "cmake_find_package") default_options = { diff --git a/internal/core/src/common/Geometry.h b/internal/core/src/common/Geometry.h index 4ef1ecfb670..b8c0ad10069 100644 --- a/internal/core/src/common/Geometry.h +++ b/internal/core/src/common/Geometry.h @@ -32,6 +32,21 @@ class Geometry { to_wkb_internal(); } + // lightweight constructor: parse wkb but **do not** copy it back to internal buffer, + // which avoids an extra malloc + memcpy. Suitable for short-lived, read-only objects + // where we only need spatial predicates (equals/intersects/…). + explicit Geometry(const void* wkb, size_t size, bool copy_wkb) { + OGRGeometry* geometry = nullptr; + OGRGeometryFactory::createFromWkb(wkb, nullptr, &geometry, size); + AssertInfo(geometry != nullptr, + "failed to construct geometry from wkb data"); + geometry_.reset(geometry); + size_ = size; + if (copy_wkb) { + to_wkb_internal(); + } + } + explicit Geometry(const char* wkt) { OGRGeometry* geometry = nullptr; OGRGeometryFactory::createFromWkt(wkt, nullptr, &geometry); diff --git a/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp b/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp index 6010cf3b044..36cf2f61659 100644 --- a/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp +++ b/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp @@ -14,6 +14,7 @@ #include "common/Geometry.h" #include "common/Types.h" #include "pb/plan.pb.h" +#include "pb/schema.pb.h" namespace milvus { namespace exec { @@ -49,8 +50,7 @@ PhyGISFunctionFilterExpr::Eval(EvalCtx& context, VectorPtr& result) { "unsupported data type: {}", expr_->column_.data_type_); if (is_index_mode_) { - // result = EvalForIndexSegment(); - PanicInfo(NotImplemented, "index for geos not implement"); + result = EvalForIndexSegment(); } else { result = EvalForDataSegment(); } @@ -143,10 +143,216 @@ PhyGISFunctionFilterExpr::EvalForDataSegment() { return res_vec; } -// VectorPtr -// PhyGISFunctionFilterExpr::EvalForIndexSegment() { -// // TODO -// } +VectorPtr +PhyGISFunctionFilterExpr::EvalForIndexSegment() { + auto real_batch_size = GetNextBatchSize(); + if (real_batch_size == 0) { + return nullptr; + } + + using Index = index::ScalarIndex; + + // Prepare shared dataset for index query (coarse candidate set by R-Tree) + auto ds = std::make_shared(); + ds->Set(milvus::index::OPERATOR_TYPE, expr_->op_); + ds->Set(milvus::index::MATCH_VALUE, expr_->geometry_.to_wkb_string()); + + /* ------------------------------------------------------------------ + * Prefetch: if coarse results are not cached yet, run a single R-Tree + * query for all index chunks and cache their coarse bitmaps. + * ------------------------------------------------------------------*/ + if (!coarse_cached_) { + // Query segment-level R-Tree index **once** since each chunk shares the same index + const Index& idx_ref = + segment_->chunk_scalar_index(field_id_, 0); + auto* idx_ptr = const_cast(&idx_ref); + + { + auto tmp = idx_ptr->Query(ds); + coarse_global_ = std::move(tmp); + } + { + auto tmp_valid = idx_ptr->IsNotNull(); + coarse_valid_global_ = std::move(tmp_valid); + } + + coarse_cached_ = true; + } + + TargetBitmap batch_result; + TargetBitmap batch_valid; + int processed_rows = 0; + + for (size_t i = current_index_chunk_; i < num_index_chunk_; ++i) { + // 1) Build and cache refined bitmap for this chunk (coarse + exact) + if (cached_index_chunk_id_ != static_cast(i)) { + // Reuse segment-level coarse cache directly + auto& coarse = coarse_global_; + auto& chunk_valid = coarse_valid_global_; + // Exact refinement + TargetBitmap refined(coarse.size()); + const bool is_sealed = segment_->type() == SegmentType::Sealed; + if (is_sealed) { + // Collect all hit row offsets from coarse bitmap + std::vector hit_offsets; + hit_offsets.reserve( + coarse.count()); // Reserve space for efficiency + for (size_t i = 0; i < coarse.size(); ++i) { + if (coarse[i]) { + hit_offsets.push_back(static_cast(i)); + } + } + + if (!hit_offsets.empty()) { + // Bulk get data for all hit rows at once + auto data_array = segment_->bulk_subscript( + field_id_, hit_offsets.data(), hit_offsets.size()); + + // Process each hit row + auto geometry_array = static_cast< + const milvus::proto::schema::GeometryArray*>( + &data_array->scalars().geometry_data()); + const auto& valid_data = data_array->valid_data(); + + for (size_t i = 0; i < hit_offsets.size(); ++i) { + const auto pos = hit_offsets[i]; + + // Check validity if available + if (!valid_data.empty() && !valid_data[i]) { + continue; + } + + const auto& wkb_data = geometry_array->data(i); + Geometry left(wkb_data.data(), wkb_data.size(), false); + bool ok = false; + + switch (expr_->op_) { + case proto::plan:: + GISFunctionFilterExpr_GISOp_Equals: + ok = left.equals(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Touches: + ok = left.touches(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Overlaps: + ok = left.overlaps(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Crosses: + ok = left.crosses(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Contains: + ok = left.contains(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Intersects: + ok = left.intersects(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Within: + ok = left.within(expr_->geometry_); + break; + default: + PanicInfo(NotImplemented, + "unknown GIS op : {}", + expr_->op_); + } + + if (ok) { + refined.set(pos); + } + } + } + } else { // Growing segment + auto span = segment_->chunk_data(field_id_, i); + + const auto start_pos = + segment_->num_rows_until_chunk(field_id_, i); + const auto chunk_rows = span.row_count(); + const auto max_local = std::min( + chunk_rows, + coarse.size() > start_pos ? coarse.size() - start_pos : 0); + + for (size_t local = 0; local < max_local; ++local) { + const size_t pos = start_pos + local; + if (!coarse[pos]) + continue; + + const auto& wkb = span[local]; + Geometry left(wkb.data(), wkb.size(), false); + bool ok = false; + switch (expr_->op_) { + case proto::plan::GISFunctionFilterExpr_GISOp_Equals: + ok = left.equals(expr_->geometry_); + break; + case proto::plan::GISFunctionFilterExpr_GISOp_Touches: + ok = left.touches(expr_->geometry_); + break; + case proto::plan::GISFunctionFilterExpr_GISOp_Overlaps: + ok = left.overlaps(expr_->geometry_); + break; + case proto::plan::GISFunctionFilterExpr_GISOp_Crosses: + ok = left.crosses(expr_->geometry_); + break; + case proto::plan::GISFunctionFilterExpr_GISOp_Contains: + ok = left.contains(expr_->geometry_); + break; + case proto::plan:: + GISFunctionFilterExpr_GISOp_Intersects: + ok = left.intersects(expr_->geometry_); + break; + case proto::plan::GISFunctionFilterExpr_GISOp_Within: + ok = left.within(expr_->geometry_); + break; + default: + PanicInfo(NotImplemented, + "unknown GIS op : {}", + expr_->op_); + } + if (ok) { + refined.set(pos); + } + } + } + + // Cache refined result for reuse by subsequent batches + cached_index_chunk_id_ = i; + cached_index_chunk_res_ = std::move(refined); + // No need to copy valid bitmap into member; use coarse_valid_cache_[i] directly later + } + + // 2) Append this chunk's cached results into current batch window + const auto& chunk_valid_ref = this->coarse_valid_global_; + + auto size = ProcessIndexOneChunk(batch_result, + batch_valid, + i, + cached_index_chunk_res_, + chunk_valid_ref, + processed_rows); + LOG_INFO( + "EvalForIndexSegment loop - i: {}, " + "processed_rows_before_append: {}, size_from_ProcessIndexOneChunk: " + "{}, batch_size_: {}", + i, + processed_rows, + size, + batch_size_); + if (processed_rows + size >= batch_size_) { + current_index_chunk_ = i; + current_index_chunk_pos_ = i == current_index_chunk_ + ? current_index_chunk_pos_ + size + : size; + break; + } + processed_rows += size; + } + return std::make_shared(std::move(batch_result), + std::move(batch_valid)); +} } //namespace exec } // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/expression/GISFunctionFilterExpr.h b/internal/core/src/exec/expression/GISFunctionFilterExpr.h index 28d5687b065..eac137d1a77 100644 --- a/internal/core/src/exec/expression/GISFunctionFilterExpr.h +++ b/internal/core/src/exec/expression/GISFunctionFilterExpr.h @@ -48,14 +48,25 @@ class PhyGISFunctionFilterExpr : public SegmentExpr { Eval(EvalCtx& context, VectorPtr& result) override; private: - // VectorPtr - // EvalForIndexSegment(); + VectorPtr + EvalForIndexSegment(); VectorPtr EvalForDataSegment(); private: std::shared_ptr expr_; + + /* + * Segment-level cache: run a single R-Tree Query for all index chunks to + * obtain coarse candidate bitmaps. Subsequent batches reuse these cached + * results to avoid repeated ScalarIndex::Query calls per chunk. + */ + bool coarse_cached_ = + false; // whether coarse results have been prefetched once + TargetBitmap coarse_global_; // global coarse bitmap (segment-level) + TargetBitmap + coarse_valid_global_; // global not-null bitmap (segment-level) }; } //namespace exec } // namespace milvus diff --git a/internal/core/src/index/IndexFactory.cpp b/internal/core/src/index/IndexFactory.cpp index bfdfa34a9ea..57fc1ffbc70 100644 --- a/internal/core/src/index/IndexFactory.cpp +++ b/internal/core/src/index/IndexFactory.cpp @@ -34,6 +34,7 @@ #include "index/BoolIndex.h" #include "index/InvertedIndexTantivy.h" #include "index/HybridScalarIndex.h" +#include "index/RTreeIndex.h" #include "knowhere/comp/knowhere_check.h" #include "log/Log.h" #include "pb/schema.pb.h" @@ -409,6 +410,15 @@ IndexFactory::CreateJsonIndex( } } +IndexBasePtr +IndexFactory::CreateGeometryIndex( + IndexType index_type, + const storage::FileManagerContext& file_manager_context) { + AssertInfo(index_type == RTREE_INDEX_TYPE, + "Invalid index type for geometry index"); + return std::make_unique>(file_manager_context); +} + IndexBasePtr IndexFactory::CreateScalarIndex( const CreateIndexInfo& create_index_info, @@ -437,6 +447,10 @@ IndexFactory::CreateScalarIndex( file_manager_context, create_index_info.json_cast_function); } + case DataType::GEOMETRY: { + return CreateGeometryIndex(create_index_info.index_type, + file_manager_context); + } default: PanicInfo(DataTypeInvalid, "Invalid data type:{}", data_type); } diff --git a/internal/core/src/index/IndexFactory.h b/internal/core/src/index/IndexFactory.h index 42cf952fc67..869ddd28d8f 100644 --- a/internal/core/src/index/IndexFactory.h +++ b/internal/core/src/index/IndexFactory.h @@ -116,6 +116,12 @@ class IndexFactory { storage::FileManagerContext(), const std::string& json_cast_function = UNKNOW_CAST_FUNCTION_NAME); + IndexBasePtr + CreateGeometryIndex( + IndexType index_type, + const storage::FileManagerContext& file_manager_context = + storage::FileManagerContext()); + IndexBasePtr CreateScalarIndex(const CreateIndexInfo& create_index_info, const storage::FileManagerContext& file_manager_context = diff --git a/internal/core/src/index/Meta.h b/internal/core/src/index/Meta.h index c904df12ccd..eacc501874c 100644 --- a/internal/core/src/index/Meta.h +++ b/internal/core/src/index/Meta.h @@ -46,6 +46,7 @@ constexpr const char* MARISA_TRIE_UPPER = "TRIE"; constexpr const char* INVERTED_INDEX_TYPE = "INVERTED"; constexpr const char* BITMAP_INDEX_TYPE = "BITMAP"; constexpr const char* HYBRID_INDEX_TYPE = "HYBRID"; +constexpr const char* RTREE_INDEX_TYPE = "RTREE"; constexpr const char* SCALAR_INDEX_ENGINE_VERSION = "scalar_index_engine_version"; constexpr const char* INDEX_NON_ENCODING = "index.nonEncoding"; diff --git a/internal/core/src/index/RTreeIndex.cpp b/internal/core/src/index/RTreeIndex.cpp new file mode 100644 index 00000000000..6a0d31ed8b8 --- /dev/null +++ b/internal/core/src/index/RTreeIndex.cpp @@ -0,0 +1,527 @@ +// 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 "index/RTreeIndex.h" +#include +#include +#include +#include "common/Slice.h" // for INDEX_FILE_SLICE_META and Disassemble +#include "common/EasyAssert.h" +#include "common/FieldData.h" +#include "log/Log.h" +#include "index/Utils.h" +#include "index/Meta.h" +#include "storage/LocalChunkManagerSingleton.h" +#include "pb/schema.pb.h" + +namespace milvus::index { + +constexpr const char* TMP_RTREE_INDEX_PREFIX = "/tmp/milvus/rtree-index/"; + +// helper to check suffix +static inline bool +ends_with(const std::string& value, const std::string& suffix) { + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == + 0; +} + +template +void +RTreeIndex::InitForBuildIndex() { + auto field = + std::to_string(disk_file_manager_->GetFieldDataMeta().field_id); + auto prefix = disk_file_manager_->GetIndexIdentifier(); + path_ = std::string(TMP_RTREE_INDEX_PREFIX) + prefix; + boost::filesystem::create_directories(path_); + + std::string index_file_path = path_ + "/index_file"; // base path (no ext) + + if (boost::filesystem::exists(index_file_path + ".dat") || + boost::filesystem::exists(index_file_path + ".idx")) { + PanicInfo( + IndexBuildError, "build rtree index temp dir:{} not empty", path_); + } + wrapper_ = std::make_shared(index_file_path, true); +} + +template +RTreeIndex::RTreeIndex(const storage::FileManagerContext& ctx) + : ScalarIndex(RTREE_INDEX_TYPE), + schema_(ctx.fieldDataMeta.field_schema) { + mem_file_manager_ = std::make_shared(ctx); + disk_file_manager_ = std::make_shared(ctx); + + if (ctx.for_loading_index) { + return; + } +} + +template +RTreeIndex::~RTreeIndex() { + // Free wrapper explicitly to ensure files not being used + wrapper_.reset(); + + // Remove temporary directory if it exists + if (!path_.empty()) { + auto local_cm = storage::LocalChunkManagerSingleton::GetInstance() + .GetChunkManager(); + if (local_cm) { + LOG_INFO("rtree index remove path:{}", path_); + local_cm->RemoveDir(path_); + } + } +} + +static std::string +GetFileName(const std::string& path) { + auto pos = path.find_last_of('/'); + return pos == std::string::npos ? path : path.substr(pos + 1); +} + +// Loading existing R-Tree index +// The config must contain "index_files" -> vector +// Remote index objects will be downloaded to local disk via DiskFileManager, +// then RTreeIndexWrapper will load them. +template +void +RTreeIndex::Load(milvus::tracer::TraceContext ctx, const Config& config) { + LOG_DEBUG("Load RTreeIndex with config {}", config.dump()); + + auto index_files_opt = + GetValueFromConfig>(config, "index_files"); + AssertInfo(index_files_opt.has_value(), + "index file paths are empty when loading R-Tree index"); + + auto files = index_files_opt.value(); + + // 1. Extract and load null_offset file(s) if present + { + auto find_file = [&](const std::string& target) -> auto { + return std::find_if( + files.begin(), files.end(), [&](const std::string& filename) { + return GetFileName(filename) == target; + }); + }; + + auto fill_null_offsets = [&](const uint8_t* data, int64_t size) { + folly::SharedMutexWritePriority::WriteHolder lock(mutex_); + null_offset_.resize((size_t)size / sizeof(size_t)); + memcpy(null_offset_.data(), data, (size_t)size); + }; + + std::vector null_offset_files; + if (auto it = find_file(INDEX_FILE_SLICE_META); it != files.end()) { + // sliced case: collect all parts with prefix index_null_offset + null_offset_files.push_back(*it); + for (auto& f : files) { + auto filename = GetFileName(f); + static const std::string kName = "index_null_offset"; + if (filename.size() >= kName.size() && + filename.substr(0, kName.size()) == kName) { + null_offset_files.push_back(f); + } + } + if (!null_offset_files.empty()) { + auto index_datas = + mem_file_manager_->LoadIndexToMemory(null_offset_files); + auto compacted = CompactIndexDatas(index_datas); + auto codecs = std::move(compacted.at("index_null_offset")); + for (auto&& codec : codecs.codecs_) { + fill_null_offsets(codec->PayloadData(), + codec->PayloadSize()); + } + } + } else if (auto it = find_file("index_null_offset"); + it != files.end()) { + null_offset_files.push_back(*it); + files.erase(it); + auto index_datas = mem_file_manager_->LoadIndexToMemory( + {*null_offset_files.begin()}); + auto null_data = std::move(index_datas.at("index_null_offset")); + fill_null_offsets(null_data->PayloadData(), + null_data->PayloadSize()); + } + + // remove loaded null_offset files from files list + if (!null_offset_files.empty()) { + files.erase(std::remove_if( + files.begin(), + files.end(), + [&](const std::string& f) { + return std::find(null_offset_files.begin(), + null_offset_files.end(), + f) != null_offset_files.end(); + }), + files.end()); + } + } + + // 2. Ensure each file has full remote path. If only filename provided, prepend remote prefix. + for (auto& f : files) { + boost::filesystem::path p(f); + if (!p.has_parent_path()) { + auto remote_prefix = disk_file_manager_->GetRemoteIndexPrefix(); + f = remote_prefix + "/" + f; + } + } + + // 3. Cache remote index files to local disk. + disk_file_manager_->CacheIndexToDisk(files); + + // 4. Determine local base path (without extension) for RTreeIndexWrapper. + auto local_paths = disk_file_manager_->GetLocalFilePaths(); + AssertInfo(!local_paths.empty(), + "RTreeIndex local files are empty after caching to disk"); + + // Pick a .dat or .idx file explicitly; avoid meta or others. + std::string base_path; + for (const auto& p : local_paths) { + if (ends_with(p, ".bgi")) { + base_path = p.substr(0, p.size() - 4); + break; + } + } + // Fallback: if not found, try meta json + if (base_path.empty()) { + for (const auto& p : local_paths) { + if (ends_with(p, ".meta.json")) { + base_path = + p.substr(0, p.size() - std::string(".meta.json").size()); + break; + } + } + } + // Final fallback: use the first path as-is + if (base_path.empty()) { + base_path = local_paths.front(); + } + path_ = base_path; + + // 5. Instantiate wrapper and load. + wrapper_ = + std::make_shared(path_, /*is_build_mode=*/false); + wrapper_->load(); + + total_num_rows_ = + wrapper_->count() + static_cast(null_offset_.size()); + is_built_ = true; + + LOG_INFO( + "Loaded R-Tree index from {} with {} rows", path_, total_num_rows_); +} + +template +void +RTreeIndex::Build(const Config& config) { + auto insert_files = + GetValueFromConfig>(config, "insert_files"); + AssertInfo(insert_files.has_value(), + "insert_files were empty for building RTree index"); + InitForBuildIndex(); + + // load raw WKB data into memory + auto field_datas = + mem_file_manager_->CacheRawDataToMemory(insert_files.value()); + BuildWithFieldData(field_datas); + // after build, mark built + total_num_rows_ = + wrapper_->count() + static_cast(null_offset_.size()); + is_built_ = true; +} + +template +void +RTreeIndex::BuildWithFieldData( + const std::vector& field_datas) { + // Default to bulk load for build performance + // If needed, we can wire a config switch later to disable it. + bool use_bulk_load = true; + if (use_bulk_load) { + // Single pass: collect null offsets locally and compute total rows + int64_t total_rows = 0; + if (schema_.nullable()) { + std::vector local_nulls; + int64_t global_offset = 0; + for (const auto& fd : field_datas) { + const auto n = fd->get_num_rows(); + for (int64_t i = 0; i < n; ++i) { + if (!fd->is_valid(i)) { + local_nulls.push_back( + static_cast(global_offset)); + } + ++global_offset; + } + total_rows += n; + } + if (!local_nulls.empty()) { + folly::SharedMutexWritePriority::WriteHolder lock(mutex_); + null_offset_.reserve(null_offset_.size() + local_nulls.size()); + null_offset_.insert( + null_offset_.end(), local_nulls.begin(), local_nulls.end()); + } + } else { + for (const auto& fd : field_datas) { + total_rows += fd->get_num_rows(); + } + } + // bulk load non-null geometries + wrapper_->bulk_load_from_field_data(field_datas, schema_.nullable()); + total_num_rows_ = total_rows; + is_built_ = true; + return; + } +} + +template +void +RTreeIndex::finish() { + if (wrapper_) { + LOG_INFO("rtree index finish"); + wrapper_->finish(); + } +} + +template +IndexStatsPtr +RTreeIndex::Upload(const Config& config) { + // 1. Ensure all buffered data flushed to disk + finish(); + + // 2. Walk temp dir and register files to DiskFileManager + boost::filesystem::path dir(path_); + boost::filesystem::directory_iterator end_iter; + + for (boost::filesystem::directory_iterator it(dir); it != end_iter; ++it) { + if (boost::filesystem::is_directory(*it)) { + LOG_WARN("{} is a directory, skip", it->path().string()); + continue; + } + + AssertInfo(disk_file_manager_->AddFile(it->path().string()), + "failed to add index file: {}", + it->path().string()); + } + + // 3. Collect remote paths to size mapping + auto remote_paths_to_size = disk_file_manager_->GetRemotePathsToFileSize(); + + // 4. Serialize and register in-memory null_offset if any + auto binary_set = Serialize(config); + mem_file_manager_->AddFile(binary_set); + auto remote_mem_path_to_size = + mem_file_manager_->GetRemotePathsToFileSize(); + + // 5. Assemble IndexStats result + std::vector index_files; + index_files.reserve(remote_paths_to_size.size() + + remote_mem_path_to_size.size()); + for (auto& kv : remote_paths_to_size) { + index_files.emplace_back(kv.first, kv.second); + } + for (auto& kv : remote_mem_path_to_size) { + index_files.emplace_back(kv.first, kv.second); + } + + int64_t mem_size = mem_file_manager_->GetAddedTotalMemSize(); + int64_t file_size = disk_file_manager_->GetAddedTotalFileSize(); + + return IndexStats::New(mem_size + file_size, std::move(index_files)); +} + +template +BinarySet +RTreeIndex::Serialize(const Config& config) { + folly::SharedMutexWritePriority::ReadHolder lock(mutex_); + auto bytes = null_offset_.size() * sizeof(size_t); + BinarySet res_set; + if (bytes > 0) { + std::shared_ptr buf(new uint8_t[bytes]); + std::memcpy(buf.get(), null_offset_.data(), bytes); + res_set.Append("index_null_offset", buf, bytes); + } + milvus::Disassemble(res_set); + return res_set; +} + +template +void +RTreeIndex::Load(const BinarySet& binary_set, const Config& config) { + PanicInfo(ErrorCode::NotImplemented, + "Load(BinarySet) is not yet supported for RTreeIndex"); +} + +template +void +RTreeIndex::Build(size_t n, const T* values, const bool* valid_data) { + // Generic Build by value array is not required for RTree at the moment. + PanicInfo(ErrorCode::NotImplemented, + "Build(size_t, values, valid) not supported for RTreeIndex"); +} + +template +const TargetBitmap +RTreeIndex::In(size_t n, const T* values) { + PanicInfo(ErrorCode::NotImplemented, "In() not supported for RTreeIndex"); + return {}; +} + +template +const TargetBitmap +RTreeIndex::IsNull() { + int64_t count = Count(); + TargetBitmap bitset(count); + folly::SharedMutexWritePriority::ReadHolder lock(mutex_); + auto end = std::lower_bound( + null_offset_.begin(), null_offset_.end(), static_cast(count)); + for (auto it = null_offset_.begin(); it != end; ++it) { + bitset.set(*it); + } + return bitset; +} + +template +const TargetBitmap +RTreeIndex::IsNotNull() { + int64_t count = Count(); + TargetBitmap bitset(count, true); + folly::SharedMutexWritePriority::ReadHolder lock(mutex_); + auto end = std::lower_bound( + null_offset_.begin(), null_offset_.end(), static_cast(count)); + for (auto it = null_offset_.begin(); it != end; ++it) { + bitset.reset(*it); + } + return bitset; +} + +template +const TargetBitmap +RTreeIndex::InApplyFilter(size_t n, + const T* values, + const std::function& filter) { + PanicInfo(ErrorCode::NotImplemented, + "InApplyFilter() not supported for RTreeIndex"); + return {}; +} + +template +void +RTreeIndex::InApplyCallback(size_t n, + const T* values, + const std::function& callback) { + PanicInfo(ErrorCode::NotImplemented, + "InApplyCallback() not supported for RTreeIndex"); +} + +template +const TargetBitmap +RTreeIndex::NotIn(size_t n, const T* values) { + PanicInfo(ErrorCode::NotImplemented, + "NotIn() not supported for RTreeIndex"); + return {}; +} + +template +const TargetBitmap +RTreeIndex::Range(T value, OpType op) { + PanicInfo(ErrorCode::NotImplemented, + "Range(value, op) not supported for RTreeIndex"); + return {}; +} + +template +const TargetBitmap +RTreeIndex::Range(T lower_bound_value, + bool lb_inclusive, + T upper_bound_value, + bool ub_inclusive) { + PanicInfo(ErrorCode::NotImplemented, + "Range(lower, upper) not supported for RTreeIndex"); + return {}; +} + +template +void +RTreeIndex::QueryCandidates(proto::plan::GISFunctionFilterExpr_GISOp op, + const std::string& query_geom_wkb, + std::vector& candidate_offsets) { + AssertInfo(wrapper_ != nullptr, "R-Tree index wrapper is null"); + OGRGeometry* geom = nullptr; + // Create OGRGeometry from WKB bytes + OGRGeometryFactory::createFromWkb( + reinterpret_cast(query_geom_wkb.data()), + nullptr, + &geom, + query_geom_wkb.size()); + AssertInfo(geom != nullptr, "invalid query geometry wkb"); + std::unique_ptr holder(geom); + wrapper_->query_candidates(op, *geom, candidate_offsets); +} + +template +const TargetBitmap +RTreeIndex::Query(const DatasetPtr& dataset) { + AssertInfo(schema_.data_type() == proto::schema::DataType::Geometry, + "RTreeIndex can only be queried on geometry field"); + auto op = + dataset->Get(OPERATOR_TYPE); + // Query geometry WKB passed via MATCH_VALUE as std::string + auto wkb = dataset->Get(MATCH_VALUE); + + // 1) Coarse candidates by R-Tree on MBR + std::vector candidate_offsets; + QueryCandidates(op, wkb, candidate_offsets); + + // 2) Build initial bitmap from candidates + TargetBitmap res(this->Count()); + for (auto off : candidate_offsets) { + if (off >= 0 && off < res.size()) { + res.set(off); + } + } + + return res; +} + +// ------------------------------------------------------------------ +// BuildWithRawDataForUT – real implementation for unit-test scenarios +// ------------------------------------------------------------------ + +template +void +RTreeIndex::BuildWithRawDataForUT(size_t n, + const void* values, + const Config& config) { + // In UT we directly receive an array of std::string (WKB) with length n. + const std::string* wkb_array = reinterpret_cast(values); + + // Guard: n should represent number of strings not raw bytes + AssertInfo(n > 0, "BuildWithRawDataForUT expects element count > 0"); + LOG_WARN("BuildWithRawDataForUT:{}", n); + this->InitForBuildIndex(); + + int64_t offset = 0; + for (size_t i = 0; i < n; ++i) { + const auto& wkb = wkb_array[i]; + const uint8_t* data_ptr = reinterpret_cast(wkb.data()); + this->wrapper_->add_geometry(data_ptr, wkb.size(), offset++); + } + this->finish(); + LOG_WARN("BuildWithRawDataForUT finish"); + this->total_num_rows_ = offset; + LOG_WARN("BuildWithRawDataForUT total_num_rows_:{}", this->total_num_rows_); + this->is_built_ = true; +} + +// Explicit template instantiation for std::string as we only support string field for now. +template class RTreeIndex; + +} // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/RTreeIndex.h b/internal/core/src/index/RTreeIndex.h new file mode 100644 index 00000000000..d8fbec50bf3 --- /dev/null +++ b/internal/core/src/index/RTreeIndex.h @@ -0,0 +1,174 @@ +// 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 "storage/FileManager.h" +#include "storage/DiskFileManagerImpl.h" +#include "storage/MemFileManagerImpl.h" +#include "index/RTreeIndexWrapper.h" +#include "index/ScalarIndex.h" +#include "pb/plan.pb.h" + +namespace milvus::index { + +using RTreeIndexWrapper = milvus::index::RTreeIndexWrapper; + +template +class RTreeIndex : public ScalarIndex { + public: + using MemFileManager = storage::MemFileManagerImpl; + using MemFileManagerPtr = std::shared_ptr; + using DiskFileManager = storage::DiskFileManagerImpl; + using DiskFileManagerPtr = std::shared_ptr; + + RTreeIndex() : ScalarIndex(RTREE_INDEX_TYPE) { + } + + explicit RTreeIndex(const storage::FileManagerContext& ctx); + + ~RTreeIndex(); + + void + InitForBuildIndex(); + + void + Load(milvus::tracer::TraceContext ctx, const Config& config = {}) override; + + // Load index from an already assembled BinarySet (not used by RTree yet) + void + Load(const BinarySet& binary_set, const Config& config = {}) override; + + ScalarIndexType + GetIndexType() const override { + return ScalarIndexType::RTREE; + } + + void + Build(const Config& config = {}) override; + + // Build index directly from in-memory value array (required by ScalarIndex) + void + Build(size_t n, const T* values, const bool* valid_data = nullptr) override; + + int64_t + Count() override { + if (is_built_) { + return total_num_rows_; + } + return wrapper_ ? wrapper_->count() + + static_cast(null_offset_.size()) + : 0; + } + + // BuildWithRawDataForUT should be only used in ut. Only string is supported. + void + BuildWithRawDataForUT(size_t n, + const void* values, + const Config& config = {}) override; + + BinarySet + Serialize(const Config& config) override; + + IndexStatsPtr + Upload(const Config& config = {}) override; + + const TargetBitmap + In(size_t n, const T* values) override; + + const TargetBitmap + IsNull() override; + + const TargetBitmap + IsNotNull() override; + + const TargetBitmap + InApplyFilter( + size_t n, + const T* values, + const std::function& filter) override; + + void + InApplyCallback( + size_t n, + const T* values, + const std::function& callback) override; + + const TargetBitmap + NotIn(size_t n, const T* values) override; + + const TargetBitmap + Range(T value, OpType op) override; + + const TargetBitmap + Range(T lower_bound_value, + bool lb_inclusive, + T upper_bound_value, + bool ub_inclusive) override; + + const bool + HasRawData() const override { + return false; + } + + std::optional + Reverse_Lookup(size_t offset) const override { + PanicInfo(ErrorCode::NotImplemented, + "Reverse_Lookup should not be handled by R-Tree index"); + } + + int64_t + Size() override { + return Count(); + } + + // GIS-specific query methods + /** + * @brief Query candidates based on spatial operation + * @param op Spatial operation type + * @param query_geom Query geometry in WKB format + * @param candidate_offsets Output vector of candidate row offsets + */ + void + QueryCandidates(proto::plan::GISFunctionFilterExpr_GISOp op, + const std::string& query_geom_wkb, + std::vector& candidate_offsets); + + const TargetBitmap + Query(const DatasetPtr& dataset) override; + + void + BuildWithFieldData(const std::vector& datas) override; + + protected: + void + finish(); + + protected: + std::shared_ptr wrapper_; + std::string path_; + proto::schema::FieldSchema schema_; + + MemFileManagerPtr mem_file_manager_; + DiskFileManagerPtr disk_file_manager_; + + // Index state + bool is_built_ = false; + int64_t total_num_rows_ = 0; + + // Track null rows to support IsNull/IsNotNull just like other scalar indexes + folly::SharedMutexWritePriority mutex_{}; + std::vector null_offset_; +}; +} // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/RTreeIndexSerialization.h b/internal/core/src/index/RTreeIndexSerialization.h new file mode 100644 index 00000000000..c0f36f0be76 --- /dev/null +++ b/internal/core/src/index/RTreeIndexSerialization.h @@ -0,0 +1,147 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +class RTreeSerializer { + public: + template + static bool + saveBinary(const RTreeType& tree, const std::string& filename) { + try { + std::ofstream ofs(filename, std::ios::binary); + if (!ofs.is_open()) { + std::cerr << "Cannot open file for writing: " << filename + << std::endl; + return false; + } + + boost::archive::binary_oarchive oa(ofs); + oa << tree; + + ofs.close(); + return true; + } catch (const std::exception& e) { + std::cerr << "Serialization error: " << e.what() << std::endl; + return false; + } + } + + template + static bool + loadBinary(RTreeType& tree, const std::string& filename) { + try { + std::ifstream ifs(filename, std::ios::binary); + if (!ifs.is_open()) { + std::cerr << "Cannot open file for reading: " << filename + << std::endl; + return false; + } + + boost::archive::binary_iarchive ia(ifs); + ia >> tree; + + ifs.close(); + return true; + } catch (const std::exception& e) { + std::cerr << "Deserialization error: " << e.what() << std::endl; + return false; + } + } + + template + static bool + saveText(const RTreeType& tree, const std::string& filename) { + try { + std::ofstream ofs(filename); + if (!ofs.is_open()) { + std::cerr << "Cannot open file for writing: " << filename + << std::endl; + return false; + } + + boost::archive::text_oarchive oa(ofs); + oa << tree; + + ofs.close(); + return true; + } catch (const std::exception& e) { + std::cerr << "Serialization error: " << e.what() << std::endl; + return false; + } + } + + template + static bool + loadText(RTreeType& tree, const std::string& filename) { + try { + std::ifstream ifs(filename); + if (!ifs.is_open()) { + std::cerr << "Cannot open file for reading: " << filename + << std::endl; + return false; + } + + boost::archive::text_iarchive ia(ifs); + ia >> tree; + + ifs.close(); + return true; + } catch (const std::exception& e) { + std::cerr << "Deserialization error: " << e.what() << std::endl; + return false; + } + } + + template + static std::string + serializeToString(const RTreeType& tree) { + std::ostringstream oss; + boost::archive::binary_oarchive oa(oss); + oa << tree; + return oss.str(); + } + + template + static bool + deserializeFromString(RTreeType& tree, const std::string& data) { + try { + std::istringstream iss(data); + boost::archive::binary_iarchive ia(iss); + ia >> tree; + return true; + } catch (const std::exception& e) { + std::cerr << "Deserialization error: " << e.what() << std::endl; + return false; + } + } +}; diff --git a/internal/core/src/index/RTreeIndexWrapper.cpp b/internal/core/src/index/RTreeIndexWrapper.cpp new file mode 100644 index 00000000000..ec00790f9aa --- /dev/null +++ b/internal/core/src/index/RTreeIndexWrapper.cpp @@ -0,0 +1,251 @@ +// 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 "RTreeIndexWrapper.h" +#include "RTreeIndexSerialization.h" +#include "common/EasyAssert.h" +#include "log/Log.h" +#include "pb/plan.pb.h" +#include +#include +#include +#include +#include +#include "common/FieldDataInterface.h" + +namespace milvus::index { + +RTreeIndexWrapper::RTreeIndexWrapper(std::string& path, bool is_build_mode) + : index_path_(path), is_build_mode_(is_build_mode) { + if (is_build_mode_) { + std::filesystem::path dir_path = + std::filesystem::path(path).parent_path(); + if (!dir_path.empty()) { + std::filesystem::create_directories(dir_path); + } + // Start with an empty rtree for dynamic insertions + rtree_ = RTree(); + } +} + +RTreeIndexWrapper::~RTreeIndexWrapper() = default; + +void +RTreeIndexWrapper::add_geometry(const uint8_t* wkb_data, + size_t len, + int64_t row_offset) { + // Acquire write lock to protect rtree_ modification + folly::SharedMutexWritePriority::WriteHolder lock(rtree_mutex_); + + AssertInfo(is_build_mode_, "Cannot add geometry in load mode"); + std::unique_lock guard(rtree_mutex_); + + // Parse WKB data to OGR geometry + OGRGeometry* geom = nullptr; + OGRErr err = + OGRGeometryFactory::createFromWkb(wkb_data, nullptr, &geom, len); + + if (err != OGRERR_NONE || geom == nullptr) { + LOG_ERROR("Failed to parse WKB data for row {}", row_offset); + return; + } + + // Get bounding box + double minX, minY, maxX, maxY; + get_bounding_box(geom, minX, minY, maxX, maxY); + + // Create Boost box and insert + Box box(Point(minX, minY), Point(maxX, maxY)); + Value val(box, row_offset); + values_.push_back(val); + rtree_.insert(val); + + // Clean up + OGRGeometryFactory::destroyGeometry(geom); +} + +// No IDataStream; bulk-load implemented directly for Boost R-tree + +void +RTreeIndexWrapper::bulk_load_from_field_data( + const std::vector>& field_datas, + bool nullable) { + // Acquire write lock to protect rtree_ creation and modification + folly::SharedMutexWritePriority::WriteHolder lock(rtree_mutex_); + + AssertInfo(is_build_mode_, "Cannot bulk load in load mode"); + std::unique_lock guard(rtree_mutex_); + std::vector local_values; + local_values.reserve(1024); + int64_t absolute_offset = 0; + for (const auto& fd : field_datas) { + const auto n = fd->get_num_rows(); + for (int64_t i = 0; i < n; ++i, ++absolute_offset) { + const bool is_nullable_effective = nullable || fd->IsNullable(); + if (is_nullable_effective && !fd->is_valid(i)) { + continue; + } + const auto* wkb_str = + static_cast(fd->RawValue(i)); + if (wkb_str == nullptr || wkb_str->empty()) { + continue; + } + OGRGeometry* geom = nullptr; + auto err = OGRGeometryFactory::createFromWkb( + reinterpret_cast(wkb_str->data()), + nullptr, + &geom, + wkb_str->size()); + if (err != OGRERR_NONE || geom == nullptr) { + continue; + } + OGREnvelope env; + geom->getEnvelope(&env); + OGRGeometryFactory::destroyGeometry(geom); + Box box(Point(env.MinX, env.MinY), Point(env.MaxX, env.MaxY)); + local_values.emplace_back(box, absolute_offset); + } + } + values_.swap(local_values); + rtree_ = RTree(values_.begin(), values_.end()); + LOG_INFO("R-Tree bulk load (Boost) completed with {} entries", + values_.size()); +} + +void +RTreeIndexWrapper::finish() { + // Acquire write lock to protect rtree_ modification and cleanup + folly::SharedMutexWritePriority::WriteHolder lock(rtree_mutex_); + + // Guard against repeated invocations which could otherwise attempt to + // release resources multiple times (e.g. BuildWithRawDataForUT() calls + // finish(), and Upload() may call it again). + std::unique_lock guard(rtree_mutex_); + if (finished_) { + LOG_DEBUG("RTreeIndexWrapper::finish() called more than once, skip."); + return; + } + + AssertInfo(is_build_mode_, "Cannot finish in load mode"); + + // Persist to disk: write meta and binary data file + try { + // Write binary rtree data + RTreeSerializer::saveBinary(rtree_, index_path_ + ".bgi"); + + // Write meta json + nlohmann::json meta; + // index/leaf capacities are not used in Boost implementation + meta["dimension"] = dimension_; + meta["count"] = static_cast(values_.size()); + + std::ofstream ofs(index_path_ + ".meta.json", std::ios::trunc); + ofs << meta.dump(); + ofs.close(); + LOG_INFO("R-Tree meta written: {}.meta.json", index_path_); + } catch (const std::exception& e) { + LOG_WARN("Failed to write R-Tree files: {}", e.what()); + } + + finished_ = true; + + LOG_INFO("R-Tree index (Boost) finished building and saved to {}", + index_path_); +} + +void +RTreeIndexWrapper::load() { + // Acquire write lock to protect rtree_ initialization during loading + folly::SharedMutexWritePriority::WriteHolder lock(rtree_mutex_); + + AssertInfo(!is_build_mode_, "Cannot load in build mode"); + + std::unique_lock guard(rtree_mutex_); + try { + // Read meta (optional) + try { + std::ifstream ifs(index_path_ + ".meta.json"); + if (ifs.good()) { + auto meta = nlohmann::json::parse(ifs); + // index/leaf capacities are ignored for Boost implementation + if (meta.contains("dimension")) + dimension_ = meta["dimension"].get(); + } + } catch (const std::exception& e) { + LOG_WARN("Failed to read meta json: {}", e.what()); + } + + // Read binary data + RTreeSerializer::loadBinary(rtree_, index_path_ + ".bgi"); + + LOG_INFO("R-Tree index (Boost) loaded from {}", index_path_); + } catch (const std::exception& e) { + PanicInfo(ErrorCode::UnexpectedError, + fmt::format("Failed to load R-Tree index from {}: {}", + index_path_, + e.what())); + } +} + +void +RTreeIndexWrapper::query_candidates(proto::plan::GISFunctionFilterExpr_GISOp op, + const OGRGeometry& query_geom, + std::vector& candidate_offsets) { + candidate_offsets.clear(); + + // Get bounding box of query geometry + double minX, minY, maxX, maxY; + get_bounding_box(&query_geom, minX, minY, maxX, maxY); + + // Create query box + Box query_box(Point(minX, minY), Point(maxX, maxY)); + + // Perform coarse intersection query + std::vector results; + { + std::shared_lock guard(rtree_mutex_); + rtree_.query(boost::geometry::index::intersects(query_box), + std::back_inserter(results)); + } + candidate_offsets.reserve(results.size()); + for (const auto& v : results) { + candidate_offsets.push_back(v.second); + } + + LOG_DEBUG("R-Tree query returned {} candidates for operation {}", + candidate_offsets.size(), + static_cast(op)); +} + +void +RTreeIndexWrapper::get_bounding_box(const OGRGeometry* geom, + double& minX, + double& minY, + double& maxX, + double& maxY) { + AssertInfo(geom != nullptr, "Geometry is null"); + + OGREnvelope env; + geom->getEnvelope(&env); + + minX = env.MinX; + minY = env.MinY; + maxX = env.MaxX; + maxY = env.MaxY; +} + +int64_t +RTreeIndexWrapper::count() const { + return static_cast(rtree_.size()); +} + +// index/leaf capacity setters removed; not applicable for Boost rtree +} // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/RTreeIndexWrapper.h b/internal/core/src/index/RTreeIndexWrapper.h new file mode 100644 index 00000000000..f6240f52b07 --- /dev/null +++ b/internal/core/src/index/RTreeIndexWrapper.h @@ -0,0 +1,144 @@ +// 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 +#include +#include +#include "ogr_geometry.h" +#include "pb/plan.pb.h" +#include + +// Forward declaration to avoid pulling heavy field data headers here +namespace milvus { +class FieldDataBase; +} + +namespace milvus::index { + +namespace bg = boost::geometry; +namespace bgi = boost::geometry::index; + +/** + * @brief Wrapper class for boost R-Tree functionality + * + * This class provides a simplified interface to boost library, + * handling the creation, management, and querying of R-Tree spatial indexes + * for geometric data in Milvus. + */ +class RTreeIndexWrapper { + public: + /** + * @brief Constructor for RTreeIndexWrapper + * @param path Path for storing index files + * @param is_build_mode Whether this is for building new index or loading existing one + */ + explicit RTreeIndexWrapper(std::string& path, bool is_build_mode); + + /** + * @brief Destructor + */ + ~RTreeIndexWrapper(); + + /** + * @brief Add a geometry (WKB) for dynamic insertion build (UT helper) + */ + void + add_geometry(const uint8_t* wkb_data, size_t len, int64_t row_offset); + + /** + * @brief Bulk load geometries from field data (WKB strings) into a new R-Tree. + * This API will create the R-Tree via createAndBulkLoadNewRTree internally. + * @param field_datas Vector of field data blocks containing WKB strings + * @param nullable Whether the field allows nulls (null rows are skipped but offset still advances) + */ + void + bulk_load_from_field_data( + const std::vector>& + field_datas, + bool nullable); + + /** + * @brief Finish building the index and flush to disk + */ + void + finish(); + + /** + * @brief Load existing index from disk + */ + void + load(); + + /** + * @brief Query candidates based on spatial operation + * @param op Spatial operation type + * @param query_geom Query geometry + * @param candidate_offsets Output vector of candidate row offsets + */ + void + query_candidates(proto::plan::GISFunctionFilterExpr_GISOp op, + const OGRGeometry& query_geom, + std::vector& candidate_offsets); + + /** + * @brief Get the total number of geometries in the index + * @return Number of geometries + */ + int64_t + count() const; + + // Boost rtree does not use index/leaf capacities; keep only fill factor for + // compatibility (no-op currently) + + private: + /** + * @brief Get bounding box from OGR geometry + * @param geom Input geometry + * @param minX Output minimum X coordinate + * @param minY Output minimum Y coordinate + * @param maxX Output maximum X coordinate + * @param maxY Output maximum Y coordinate + */ + void + get_bounding_box(const OGRGeometry* geom, + double& minX, + double& minY, + double& maxX, + double& maxY); + + private: + // Boost.Geometry types and in-memory structures + using Point = bg::model::point>; + using Box = bg::model::box; + using Value = std::pair; // (MBR, row_offset) + using RTree = bgi::rtree>; + + RTree rtree_{}; + std::vector values_; + std::string index_path_; + bool is_build_mode_; + + // Flag to guard against repeated invocations which could otherwise attempt to release resources multiple times (e.g. BuildWithRawDataForUT() calls finish(), and Upload() may call it again). + bool finished_ = false; + + // Serialize access to rtree_ + mutable std::shared_mutex rtree_mutex_; + + // R-Tree parameters + uint32_t dimension_ = 2; +}; + +} // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/ScalarIndex.h b/internal/core/src/index/ScalarIndex.h index 1314de269cd..62d5c153901 100644 --- a/internal/core/src/index/ScalarIndex.h +++ b/internal/core/src/index/ScalarIndex.h @@ -36,6 +36,7 @@ enum class ScalarIndexType { MARISA, INVERTED, HYBRID, + RTREE, }; inline std::string @@ -53,6 +54,8 @@ ToString(ScalarIndexType type) { return "INVERTED"; case ScalarIndexType::HYBRID: return "HYBRID"; + case ScalarIndexType::RTREE: + return "RTREE"; default: return "UNKNOWN"; } diff --git a/internal/core/src/indexbuilder/IndexFactory.h b/internal/core/src/indexbuilder/IndexFactory.h index 97ecffa3290..8552d394422 100644 --- a/internal/core/src/indexbuilder/IndexFactory.h +++ b/internal/core/src/indexbuilder/IndexFactory.h @@ -62,6 +62,7 @@ class IndexFactory { case DataType::STRING: case DataType::ARRAY: case DataType::JSON: + case DataType::GEOMETRY: return CreateScalarIndex(type, config, context); case DataType::VECTOR_FLOAT: diff --git a/internal/core/unittest/CMakeLists.txt b/internal/core/unittest/CMakeLists.txt index 0aef16e660d..f02d2c34734 100644 --- a/internal/core/unittest/CMakeLists.txt +++ b/internal/core/unittest/CMakeLists.txt @@ -95,6 +95,8 @@ set(MILVUS_TEST_FILES test_json_index.cpp test_json_key_stats_index.cpp test_thread_pool.cpp + test_rtree_index_wrapper.cpp + test_rtree_index.cpp ) if(INDEX_ENGINE STREQUAL "cardinal") diff --git a/internal/core/unittest/test_expr.cpp b/internal/core/unittest/test_expr.cpp index a0d06771a25..68f6a27af45 100644 --- a/internal/core/unittest/test_expr.cpp +++ b/internal/core/unittest/test_expr.cpp @@ -17346,7 +17346,7 @@ TEST_P(ExprTest, TestGISFunctionWithControlledData) { test_gis_operation("POLYGON((-2 -2, 2 -2, 2 2, -2 2, -2 -2))", proto::plan::GISFunctionFilterExpr_GISOp_Within, [](int i) -> bool { - // Only geometry at index 0,1 (polygon containing (0,0)) + // Only geometry at index 0,1,3 (polygon containing (0,0)) return (i % 4 == 0) || (i % 4 == 1) || (i % 4 == 3); }); diff --git a/internal/core/unittest/test_rtree_index.cpp b/internal/core/unittest/test_rtree_index.cpp new file mode 100644 index 00000000000..28865bbb1da --- /dev/null +++ b/internal/core/unittest/test_rtree_index.cpp @@ -0,0 +1,866 @@ +// 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 +#include +#include +#include +#include + +#include "index/RTreeIndex.h" +#include "storage/Util.h" +#include "storage/FileManager.h" +#include "common/Types.h" +#include "test_utils/TmpPath.h" +#include "pb/schema.pb.h" +#include "pb/plan.pb.h" +#include "common/Geometry.h" +#include "common/EasyAssert.h" +#include "storage/InsertData.h" +#include "storage/PayloadReader.h" +#include "storage/DiskFileManagerImpl.h" +#include "common/FieldData.h" +#include +#include +#include "segcore/SegmentGrowingImpl.h" +#include "segcore/SegmentSealedImpl.h" +#include "test_utils/DataGen.h" +#include "query/ExecPlanNodeVisitor.h" +#include "common/Consts.h" + +// Helper: create simple POINT(x,y) WKB (little-endian) +static std::string +CreatePointWKB(double x, double y) { + std::vector wkb; + // Byte order – little endian (1) + wkb.push_back(0x01); + // Geometry type – Point (1) – 32-bit little endian + uint32_t geom_type = 1; + uint8_t* type_bytes = reinterpret_cast(&geom_type); + wkb.insert(wkb.end(), type_bytes, type_bytes + sizeof(uint32_t)); + // X coordinate + uint8_t* x_bytes = reinterpret_cast(&x); + wkb.insert(wkb.end(), x_bytes, x_bytes + sizeof(double)); + // Y coordinate + uint8_t* y_bytes = reinterpret_cast(&y); + wkb.insert(wkb.end(), y_bytes, y_bytes + sizeof(double)); + return std::string(reinterpret_cast(wkb.data()), wkb.size()); +} + +// Helper: create simple WKB from WKT +static std::string +CreateWkbFromWkt(const std::string& wkt) { + return milvus::Geometry(wkt.c_str()).to_wkb_string(); +} + +// Helper: write an InsertData parquet file to "remote" storage managed by chunk_manager_ +static std::string +WriteGeometryInsertFile(const milvus::storage::ChunkManagerPtr& cm, + const milvus::storage::FieldDataMeta& field_meta, + const std::string& remote_path, + const std::vector& wkbs, + bool nullable = false, + const uint8_t* valid_bitmap = nullptr) { + auto field_data = milvus::storage::CreateFieldData( + milvus::storage::DataType::GEOMETRY, nullable); + if (nullable && valid_bitmap != nullptr) { + field_data->FillFieldData(wkbs.data(), valid_bitmap, wkbs.size()); + } else { + field_data->FillFieldData(wkbs.data(), wkbs.size()); + } + auto payload_reader = + std::make_shared(field_data); + milvus::storage::InsertData insert_data(payload_reader); + insert_data.SetFieldDataMeta(field_meta); + insert_data.SetTimestamps(0, 100); + + auto bytes = insert_data.Serialize(milvus::storage::StorageType::Remote); + std::vector buf(bytes.begin(), bytes.end()); + cm->Write(remote_path, buf.data(), buf.size()); + return remote_path; +} + +class RTreeIndexTest : public ::testing::Test { + protected: + void + SetUp() override { + temp_path_ = milvus::test::TmpPath{}; + // create storage config that writes to temp dir + storage_config_.storage_type = "local"; + storage_config_.root_path = temp_path_.get().string(); + chunk_manager_ = milvus::storage::CreateChunkManager(storage_config_); + + // prepare field & index meta – minimal info for DiskFileManagerImpl + field_meta_ = milvus::storage::FieldDataMeta{1, 1, 1, 100}; + // set geometry data type in field schema for index schema checks + field_meta_.field_schema.set_data_type( + ::milvus::proto::schema::DataType::Geometry); + index_meta_ = milvus::storage::IndexMeta{.segment_id = 1, + .field_id = 100, + .build_id = 1, + .index_version = 1}; + } + + void + TearDown() override { + // clean chunk manager files if any (TmpPath destructor will also remove) + } + + milvus::storage::StorageConfig storage_config_; + milvus::storage::ChunkManagerPtr chunk_manager_; + milvus::storage::FieldDataMeta field_meta_; + milvus::storage::IndexMeta index_meta_; + milvus::test::TmpPath temp_path_; +}; + +TEST_F(RTreeIndexTest, Build_Upload_Load) { + // ---------- Build via BuildWithRawDataForUT ---------- + milvus::storage::FileManagerContext ctx_build( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree_build(ctx_build); + + std::vector wkbs = {CreatePointWKB(1.0, 1.0), + CreatePointWKB(2.0, 2.0)}; + rtree_build.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + + ASSERT_EQ(rtree_build.Count(), 2); + + // ---------- Upload ---------- + auto stats = rtree_build.Upload({}); + ASSERT_NE(stats, nullptr); + ASSERT_GT(stats->GetIndexFiles().size(), 0); + + // ---------- Load back ---------- + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + + milvus::tracer::TraceContext trace_ctx; // empty context + rtree_load.Load(trace_ctx, cfg); + + ASSERT_EQ(rtree_load.Count(), 2); +} + +TEST_F(RTreeIndexTest, Load_WithFileNamesOnly) { + // Build & upload first + milvus::storage::FileManagerContext ctx_build( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree_build(ctx_build); + + std::vector wkbs2 = {CreatePointWKB(10.0, 10.0), + CreatePointWKB(20.0, 20.0)}; + rtree_build.BuildWithRawDataForUT(wkbs2.size(), wkbs2.data()); + + auto stats = rtree_build.Upload({}); + + // gather only filenames (strip parent path) + std::vector filenames; + for (const auto& path : stats->GetIndexFiles()) { + filenames.emplace_back( + boost::filesystem::path(path).filename().string()); + // make sure file exists in remote storage + ASSERT_TRUE(chunk_manager_->Exist(path)); + ASSERT_GT(chunk_manager_->Size(path), 0); + } + + // Load using filename only list + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = filenames; // no directory info + + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + + ASSERT_EQ(rtree_load.Count(), 2); +} + +TEST_F(RTreeIndexTest, Build_EmptyInput_ShouldThrow) { + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + std::vector empty; + EXPECT_THROW(rtree.BuildWithRawDataForUT(0, empty.data()), + milvus::SegcoreError); +} + +TEST_F(RTreeIndexTest, Build_WithInvalidWKB_Upload_Load) { + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + std::string bad = CreatePointWKB(0.0, 0.0); + bad.resize(bad.size() / 2); // truncate to make invalid + + std::vector wkbs = { + CreateWkbFromWkt("POINT(1 1)"), bad, CreateWkbFromWkt("POINT(2 2)")}; + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + + // Upload and then load back to let loader compute count from wrapper + auto stats = rtree.Upload({}); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + + // Only 2 valid points should be present + ASSERT_EQ(rtree_load.Count(), 2); +} + +TEST_F(RTreeIndexTest, Build_VariousGeometries) { + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + std::vector wkbs = { + CreateWkbFromWkt("POINT(-1.5 2.5)"), + CreateWkbFromWkt("LINESTRING(0 0,1 1,2 3)"), + CreateWkbFromWkt("POLYGON((0 0,2 0,2 2,0 2,0 0))"), + CreateWkbFromWkt("POINT(1000000 -1000000)"), + CreateWkbFromWkt("POINT(0 0)")}; + + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + ASSERT_EQ(rtree.Count(), wkbs.size()); + + auto stats = rtree.Upload({}); + ASSERT_FALSE(stats->GetIndexFiles().empty()); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + ASSERT_EQ(rtree_load.Count(), wkbs.size()); +} + +TEST_F(RTreeIndexTest, Build_ConfigAndMetaJson) { + // Prepare one insert file via storage pipeline + std::vector wkbs = {CreateWkbFromWkt("POINT(0 0)"), + CreateWkbFromWkt("POINT(1 1)")}; + auto remote_file = (temp_path_.get() / "geom.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["fillFactor"] = "0.6"; + build_cfg["indexCapacity"] = "32"; + build_cfg["leafCapacity"] = "64"; + build_cfg["rv"] = "RSTAR"; + + rtree.Build(build_cfg); + auto stats = rtree.Upload({}); + + // Cache remote index files locally + milvus::storage::DiskFileManagerImpl diskfm( + {field_meta_, index_meta_, chunk_manager_}); + auto index_files = stats->GetIndexFiles(); + diskfm.CacheIndexToDisk(index_files); + auto local_paths = diskfm.GetLocalFilePaths(); + ASSERT_FALSE(local_paths.empty()); + // Determine base path like RTreeIndex::Load + auto ends_with = [](const std::string& value, const std::string& suffix) { + return value.size() >= suffix.size() && + value.compare( + value.size() - suffix.size(), suffix.size(), suffix) == 0; + }; + + std::string base_path; + for (const auto& p : local_paths) { + if (ends_with(p, ".dat")) { + base_path = p.substr(0, p.size() - 4); + break; + } + if (ends_with(p, ".idx")) { + base_path = p.substr(0, p.size() - 4); + break; + } + } + if (base_path.empty()) { + for (const auto& p : local_paths) { + if (ends_with(p, ".meta.json")) { + base_path = + p.substr(0, p.size() - std::string(".meta.json").size()); + break; + } + } + } + if (base_path.empty()) { + base_path = local_paths.front(); + } + // Parse local meta json + std::ifstream ifs(base_path + ".meta.json"); + ASSERT_TRUE(ifs.good()); + nlohmann::json meta = nlohmann::json::parse(ifs); + ASSERT_EQ(meta["fill_factor"], "0.6"); + ASSERT_EQ(meta["index_capacity"], "32"); + ASSERT_EQ(meta["leaf_capacity"], "64"); + ASSERT_EQ(meta["dimension"], 2); +} + +TEST_F(RTreeIndexTest, Build_InvalidVariant_ShouldThrow) { + // Prepare insert + std::vector wkbs = {CreateWkbFromWkt("POINT(0 0)")}; + auto remote_file = (temp_path_.get() / "geom2.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["rv"] = "FOO"; // invalid + EXPECT_THROW(rtree.Build(build_cfg), milvus::SegcoreError); +} + +TEST_F(RTreeIndexTest, Load_OnlyIdx_OnlyDat) { + // Build and upload + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + std::vector wkbs = {CreatePointWKB(3.0, 3.0), + CreatePointWKB(4.0, 4.0)}; + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + auto stats = rtree.Upload({}); + + std::vector only_idx, only_dat; + for (const auto& p : stats->GetIndexFiles()) { + if (boost::algorithm::ends_with(p, ".idx_0")) + only_idx.push_back(p); + if (boost::algorithm::ends_with(p, ".dat_0")) + only_dat.push_back(p); + } + ASSERT_FALSE(only_idx.empty()); + ASSERT_FALSE(only_dat.empty()); + + // Load with only idx should fail + milvus::storage::FileManagerContext ctx_load1( + field_meta_, index_meta_, chunk_manager_); + ctx_load1.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load1(ctx_load1); + nlohmann::json cfg1; + cfg1["index_files"] = only_idx; + milvus::tracer::TraceContext trace_ctx; + EXPECT_ANY_THROW(rtree_load1.Load(trace_ctx, cfg1)); + + // Load with only dat should fail + milvus::storage::FileManagerContext ctx_load2( + field_meta_, index_meta_, chunk_manager_); + ctx_load2.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load2(ctx_load2); + nlohmann::json cfg2; + cfg2["index_files"] = only_dat; + EXPECT_ANY_THROW(rtree_load2.Load(trace_ctx, cfg2)); +} + +TEST_F(RTreeIndexTest, Load_OnlyMeta_ShouldThrow) { + // Build and upload + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + std::vector wkbs = {CreatePointWKB(5.0, 5.0)}; + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + auto stats = rtree.Upload({}); + + std::vector only_meta; + for (const auto& p : stats->GetIndexFiles()) { + if (boost::algorithm::ends_with(p, ".meta.json_0")) + only_meta.push_back(p); + } + ASSERT_FALSE(only_meta.empty()); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + nlohmann::json cfg; + cfg["index_files"] = only_meta; + milvus::tracer::TraceContext trace_ctx; + EXPECT_ANY_THROW(rtree_load.Load(trace_ctx, cfg)); +} + +TEST_F(RTreeIndexTest, Load_MixedFileNamesAndPaths) { + // Build and upload + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + std::vector wkbs = {CreatePointWKB(6.0, 6.0), + CreatePointWKB(7.0, 7.0)}; + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data()); + auto stats = rtree.Upload({}); + + // Use full list, but replace one with filename-only + auto mixed = stats->GetIndexFiles(); + ASSERT_FALSE(mixed.empty()); + mixed[0] = boost::filesystem::path(mixed[0]).filename().string(); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = mixed; + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + ASSERT_EQ(rtree_load.Count(), wkbs.size()); +} + +TEST_F(RTreeIndexTest, Load_NonexistentRemote_ShouldThrow) { + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + // nonexist file + nlohmann::json cfg; + cfg["index_files"] = std::vector{ + (temp_path_.get() / "does_not_exist.idx_0").string()}; + milvus::tracer::TraceContext trace_ctx; + EXPECT_THROW(rtree_load.Load(trace_ctx, cfg), milvus::SegcoreError); +} + +TEST_F(RTreeIndexTest, Build_EndToEnd_FromInsertFiles) { + // prepare remote file via InsertData serialization + std::vector wkbs = {CreateWkbFromWkt("POINT(0 0)"), + CreateWkbFromWkt("POINT(2 2)")}; + auto remote_file = (temp_path_.get() / "geom3.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["fillFactor"] = "0.8"; + build_cfg["indexCapacity"] = "50"; + build_cfg["leafCapacity"] = "50"; + build_cfg["rv"] = "RSTAR"; + rtree.Build(build_cfg); + ASSERT_EQ(rtree.Count(), wkbs.size()); + + auto stats = rtree.Upload({}); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + ASSERT_EQ(rtree_load.Count(), wkbs.size()); +} + +TEST_F(RTreeIndexTest, Build_Upload_Load_LargeDataset) { + // Generate ~10k POINT geometries + const size_t N = 10000; + std::vector wkbs; + wkbs.reserve(N); + for (size_t i = 0; i < N; ++i) { + // POINT(i i) + wkbs.emplace_back(CreateWkbFromWkt("POINT(" + std::to_string(i) + " " + + std::to_string(i) + ")")); + } + + // Write one insert file into remote storage + auto remote_file = (temp_path_.get() / "geom_large.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + + // Build from insert_files (not using BuildWithRawDataForUT) + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["fillFactor"] = "0.8"; + build_cfg["indexCapacity"] = "50"; + build_cfg["leafCapacity"] = "50"; + build_cfg["rv"] = "RSTAR"; + rtree.Build(build_cfg); + + ASSERT_EQ(rtree.Count(), static_cast(N)); + + // Upload index + auto stats = rtree.Upload({}); + ASSERT_GT(stats->GetIndexFiles().size(), 0); + + // Load index back and verify + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg_load; + cfg_load["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg_load); + + ASSERT_EQ(rtree_load.Count(), static_cast(N)); +} + +TEST_F(RTreeIndexTest, Build_BulkLoad_Nulls_And_BadWKB) { + // five geometries: + // 1. valid + // 2. valid but will be marked null + // 3. valid + // 4. will be truncated to make invalid + // 5. valid + std::vector wkbs = { + CreateWkbFromWkt("POINT(0 0)"), // valid + CreateWkbFromWkt("POINT(1 1)"), // valid + CreateWkbFromWkt("POINT(2 2)"), // valid + CreatePointWKB(3.0, 3.0), // will be truncated to make invalid + CreateWkbFromWkt("POINT(4 4)") // valid + }; + // make bad WKB: truncate the 4th geometry + wkbs[3].resize(wkbs[3].size() / 2); + + // write to remote storage file (chunk manager's root directory) + auto remote_file = (temp_path_.get() / "geom_bulk.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + + // build (default to bulk load) + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["fillFactor"] = "0.8"; + build_cfg["indexCapacity"] = "50"; + build_cfg["leafCapacity"] = "50"; + build_cfg["rv"] = "RSTAR"; + rtree.Build(build_cfg); + + // expect: 3 geometries (0, 2, 4) are valid and parsable, 1st geometry is marked null and skipped, 3rd geometry is bad WKB and skipped + ASSERT_EQ(rtree.Count(), 4); + + // upload -> load back and verify consistency + auto stats = rtree.Upload({}); + ASSERT_GT(stats->GetIndexFiles().size(), 0); + + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + ASSERT_EQ(rtree_load.Count(), 4); +} + +// The following two tests only test the coarse query (R-Tree) and not the exact query (GDAL) +TEST_F(RTreeIndexTest, Query_CoarseAndExact_Equals_Intersects_Within) { + // Build a small index in-memory (via UT API) + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + // Prepare simple geometries: two points and a square polygon + std::vector wkbs; + wkbs.emplace_back(CreateWkbFromWkt("POINT(0 0)")); // id 0 + wkbs.emplace_back(CreateWkbFromWkt("POINT(2 2)")); // id 1 + wkbs.emplace_back( + CreateWkbFromWkt("POLYGON((0 0, 0 3, 3 3, 3 0, 0 0))")); // id 2 square + + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data(), {}); + ASSERT_EQ(rtree.Count(), 3); + + // Upload and then load into a new index instance for querying + auto stats = rtree.Upload({}); + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + + // Helper to run Query + auto run_query = [&](::milvus::proto::plan::GISFunctionFilterExpr_GISOp op, + const std::string& wkt) { + auto ds = std::make_shared(); + ds->Set(milvus::index::OPERATOR_TYPE, op); + ds->Set(milvus::index::MATCH_VALUE, CreateWkbFromWkt(wkt)); + return rtree_load.Query(ds); + }; + + // Equals with same point should match id 0 only + { + auto bm = + run_query(::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Equals, + "POINT(0 0)"); + EXPECT_TRUE(bm[0]); + EXPECT_FALSE(bm[1]); + EXPECT_TRUE( + bm[2]); //This is true because POINT(0 0) is within the square (0 0, 0 3, 3 3, 3 0, 0 0) and we have not done exact spatial query yet + } + + // Intersects: square intersects point (on boundary considered intersect) + { + auto bm = run_query( + ::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Intersects, + "POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))"); + // square(0..1) intersects POINT(0,0) and POLYGON(0..3) + // but not POINT(2,2) + EXPECT_TRUE(bm[0]); // point (0,0) + EXPECT_FALSE(bm[1]); // point (2,2) + EXPECT_TRUE(bm[2]); // big polygon + } + + // Within: point within the big square + { + auto bm = + run_query(::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Within, + "POLYGON((0 0, 0 3, 3 3, 3 0, 0 0))"); + EXPECT_TRUE( + bm[0]); // (0,0) is within or on boundary considered within by GDAL Within? + // GDAL Within returns true only if strictly inside (no boundary). If boundary excluded, (0,0) may be false. + // To make assertion robust across GEOS versions, simply check big polygon within itself should be true. + auto bm_poly = + run_query(::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Within, + "POLYGON((0 0, 0 3, 3 3, 3 0, 0 0))"); + EXPECT_TRUE(bm_poly[2]); + } +} + +TEST_F(RTreeIndexTest, Query_Touches_Contains_Crosses_Overlaps) { + milvus::storage::FileManagerContext ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree(ctx); + + // Two overlapping squares and one disjoint square + std::vector wkbs; + wkbs.emplace_back( + CreateWkbFromWkt("POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))")); // id 0 + wkbs.emplace_back(CreateWkbFromWkt( + "POLYGON((1 1, 1 3, 3 3, 3 1, 1 1))")); // id 1 overlaps with 0 + wkbs.emplace_back(CreateWkbFromWkt( + "POLYGON((4 4, 4 5, 5 5, 5 4, 4 4))")); // id 2 disjoint + + rtree.BuildWithRawDataForUT(wkbs.size(), wkbs.data(), {}); + ASSERT_EQ(rtree.Count(), 3); + + // Upload and load a new instance for querying + auto stats = rtree.Upload({}); + milvus::storage::FileManagerContext ctx_load( + field_meta_, index_meta_, chunk_manager_); + ctx_load.set_for_loading_index(true); + milvus::index::RTreeIndex rtree_load(ctx_load); + nlohmann::json cfg; + cfg["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx; + rtree_load.Load(trace_ctx, cfg); + + auto run_query = [&](::milvus::proto::plan::GISFunctionFilterExpr_GISOp op, + const std::string& wkt) { + auto ds = std::make_shared(); + ds->Set(milvus::index::OPERATOR_TYPE, op); + ds->Set(milvus::index::MATCH_VALUE, CreateWkbFromWkt(wkt)); + return rtree_load.Query(ds); + }; + + // Overlaps: query polygon overlapping both 0 and 1 + { + auto bm = run_query( + ::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Overlaps, + "POLYGON((0.5 0.5, 0.5 2.5, 2.5 2.5, 2.5 0.5, 0.5 0.5))"); + EXPECT_TRUE(bm[0]); + EXPECT_TRUE(bm[1]); + EXPECT_FALSE(bm[2]); + } + + // Contains: big polygon contains small polygon + { + auto bm = run_query( + ::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Contains, + "POLYGON(( -1 -1, -1 4, 4 4, 4 -1, -1 -1))"); + EXPECT_TRUE(bm[0]); + EXPECT_TRUE(bm[1]); + EXPECT_FALSE(bm[2]); + } + + // Touches: polygon that only touches at the corner (2,2) with id1 + { + auto bm = run_query( + ::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Touches, + "POLYGON((2 2, 2 3, 3 3, 3 2, 2 2))"); + // This touches id1 at (2,2); depending on GEOS, touches excludes interior intersection + // The id0 might also touch at (2,2). We only assert at least one touch. + EXPECT_TRUE(bm[0] || bm[1]); + } + + // Crosses: a segment crossing the first polygon + { + auto bm = run_query( + ::milvus::proto::plan::GISFunctionFilterExpr_GISOp_Crosses, + "LINESTRING( -1 1, 3 1 )"); + EXPECT_TRUE(bm[0]); + } +} + +TEST_F(RTreeIndexTest, GIS_Index_Exact_Filtering) { + using namespace milvus; + using namespace milvus::query; + using namespace milvus::segcore; + + // 1) Create schema: id (INT64, primary), vector, geometry + auto schema = std::make_shared(); + auto pk_id = schema->AddDebugField("id", DataType::INT64); + auto dim = 16; + auto vec_id = schema->AddDebugField( + "vec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2); + auto geo_id = schema->AddDebugField("geo", DataType::GEOMETRY); + schema->set_primary_field_id(pk_id); + + int N = 200; + int num_iters = 1; + // 2) Promote to sealed and build/load indices for vector + geometry + auto sealed = milvus::segcore::CreateSealedSegment(schema); + // load raw field data into sealed, excluding geometry (we will load controlled geometry separately) + auto full_ds = DataGen(schema, N * num_iters); + SealedLoadFieldData(full_ds, *sealed, {geo_id.get()}); + + // Prepare controlled geometry WKBs mirroring the shapes used in growing + std::vector wkbs; + wkbs.reserve(N * num_iters); + for (int i = 0; i < N * num_iters; ++i) { + if (i % 4 == 0) { + wkbs.emplace_back(milvus::Geometry("POINT(0 0)").to_wkb_string()); + } else if (i % 4 == 1) { + wkbs.emplace_back( + milvus::Geometry("POLYGON((-1 -1,1 -1,1 1,-1 1,-1 -1))") + .to_wkb_string()); + } else if (i % 4 == 2) { + wkbs.emplace_back( + milvus::Geometry("POLYGON((10 10,20 10,20 20,10 20,10 10))") + .to_wkb_string()); + } else { + wkbs.emplace_back( + milvus::Geometry("LINESTRING(-1 0,1 0)").to_wkb_string()); + } + } + + // now load the controlled geometry data into sealed + FieldDataInfo geo_fd_info; + geo_fd_info.field_id = geo_id.get(); + geo_fd_info.row_count = N * num_iters; + auto geo_field_data = milvus::storage::CreateFieldData( + milvus::storage::DataType::GEOMETRY, /*nullable=*/false); + geo_field_data->FillFieldData(wkbs.data(), wkbs.size()); + geo_fd_info.channel->push(geo_field_data); + geo_fd_info.channel->close(); + sealed->LoadFieldData(geo_id, geo_fd_info); + + // build geometry R-Tree index files and load into sealed + // Write a single parquet for geometry to simulate build input + // wkbs already prepared above + auto remote_file = (temp_path_.get() / "rtree_e2e.parquet").string(); + WriteGeometryInsertFile(chunk_manager_, field_meta_, remote_file, wkbs); + + // build index files by invoking RTreeIndex::Build + milvus::storage::FileManagerContext fm_ctx( + field_meta_, index_meta_, chunk_manager_); + milvus::index::RTreeIndex rtree_build(fm_ctx); + nlohmann::json build_cfg; + build_cfg["insert_files"] = std::vector{remote_file}; + build_cfg["fillFactor"] = "0.8"; + build_cfg["indexCapacity"] = "50"; + build_cfg["leafCapacity"] = "50"; + build_cfg["rv"] = "RSTAR"; + rtree_build.Build(build_cfg); + auto stats = rtree_build.Upload({}); + + // load geometry index into sealed segment + milvus::segcore::LoadIndexInfo info{}; + info.collection_id = 1; + info.partition_id = 1; + info.segment_id = 1; + info.field_id = geo_id.get(); + info.field_type = DataType::GEOMETRY; + info.index_id = 1; + info.index_build_id = 1; + info.index_version = 1; + info.schema = proto::schema::FieldSchema(); + info.schema.set_data_type(proto::schema::DataType::Geometry); + // Prepare a loaded RTree index instance and assign to info.index for scalar index loading path + milvus::storage::FileManagerContext fm_ctx_load( + field_meta_, index_meta_, chunk_manager_); + fm_ctx_load.set_for_loading_index(true); + auto rtree_loaded = + std::make_unique>(fm_ctx_load); + nlohmann::json cfg_load; + cfg_load["index_files"] = stats->GetIndexFiles(); + milvus::tracer::TraceContext trace_ctx_load; + rtree_loaded->Load(trace_ctx_load, cfg_load); + info.index = std::move(rtree_loaded); + sealed->LoadIndex(info); + + // 3) Build a GIS filter expression and run exact filtering via segcore + auto test_op = [&](const std::string& wkt, + proto::plan::GISFunctionFilterExpr_GISOp op, + std::function expected) { + milvus::Geometry right(wkt.c_str()); + auto gis_expr = std::make_shared( + milvus::expr::ColumnInfo(geo_id, DataType::GEOMETRY), op, right); + auto plan = std::make_shared(DEFAULT_PLANNODE_ID, + gis_expr); + BitsetType bits = + ExecuteQueryExpr(plan, sealed.get(), N * num_iters, MAX_TIMESTAMP); + ASSERT_EQ(bits.size(), N * num_iters); + for (int i = 0; i < N * num_iters; ++i) { + EXPECT_EQ(bool(bits[i]), expected(i)) << "i=" << i; + } + }; + + // exact within: polygon around origin should include indices 0,1,3 + test_op("POLYGON((-2 -2,2 -2,2 2,-2 2,-2 -2))", + proto::plan::GISFunctionFilterExpr_GISOp_Within, + [](int i) { return (i % 4 == 0) || (i % 4 == 1) || (i % 4 == 3); }); + + // exact intersects: point (0,0) should intersect point, polygon containing it, and line through it + test_op("POINT(0 0)", + proto::plan::GISFunctionFilterExpr_GISOp_Intersects, + [](int i) { return (i % 4 == 0) || (i % 4 == 1) || (i % 4 == 3); }); + + // exact equals: only the point equals + test_op("POINT(0 0)", + proto::plan::GISFunctionFilterExpr_GISOp_Equals, + [](int i) { return (i % 4 == 0); }); +} \ No newline at end of file diff --git a/internal/core/unittest/test_rtree_index_wrapper.cpp b/internal/core/unittest/test_rtree_index_wrapper.cpp new file mode 100644 index 00000000000..1f0c193b1ca --- /dev/null +++ b/internal/core/unittest/test_rtree_index_wrapper.cpp @@ -0,0 +1,232 @@ +// 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 +#include +#include +#include "index/RTreeIndexWrapper.h" +#include "common/Types.h" +#include "gdal.h" + +class RTreeIndexWrapperTest : public ::testing::Test { + protected: + void + SetUp() override { + // Create test directory + test_dir_ = "/tmp/rtree_test"; + std::filesystem::create_directories(test_dir_); + + // Initialize GDAL + GDALAllRegister(); + } + + void + TearDown() override { + // Clean up test directory + std::filesystem::remove_all(test_dir_); + + // Clean up GDAL + GDALDestroyDriverManager(); + } + + // Helper function to create a simple point WKB + std::vector + create_point_wkb(double x, double y) { + // WKB format for a point: byte order (1) + geometry type (1) + coordinates (16 bytes) + std::vector wkb = { + 0x01, // Little endian + 0x01, + 0x00, + 0x00, + 0x00, // Point geometry type + }; + + // Add X coordinate (8 bytes, little endian double) + uint8_t* x_bytes = reinterpret_cast(&x); + wkb.insert(wkb.end(), x_bytes, x_bytes + sizeof(double)); + + // Add Y coordinate (8 bytes, little endian double) + uint8_t* y_bytes = reinterpret_cast(&y); + wkb.insert(wkb.end(), y_bytes, y_bytes + sizeof(double)); + + return wkb; + } + + // Helper function to create a simple polygon WKB + std::vector + create_polygon_wkb(const std::vector>& points) { + // WKB format for a polygon + std::vector wkb = { + 0x01, // Little endian + 0x03, + 0x00, + 0x00, + 0x00, // Polygon geometry type + 0x01, + 0x00, + 0x00, + 0x00, // 1 ring + }; + + // Add number of points in the ring + uint32_t num_points = static_cast(points.size()); + uint8_t* num_points_bytes = reinterpret_cast(&num_points); + wkb.insert( + wkb.end(), num_points_bytes, num_points_bytes + sizeof(uint32_t)); + + // Add points + for (const auto& point : points) { + double x = point.first; + double y = point.second; + + uint8_t* x_bytes = reinterpret_cast(&x); + wkb.insert(wkb.end(), x_bytes, x_bytes + sizeof(double)); + + uint8_t* y_bytes = reinterpret_cast(&y); + wkb.insert(wkb.end(), y_bytes, y_bytes + sizeof(double)); + } + + return wkb; + } + + std::string test_dir_; +}; + +TEST_F(RTreeIndexWrapperTest, TestBuildAndLoad) { + std::string index_path = test_dir_ + "/test_index"; + + // Test building index + { + milvus::index::RTreeIndexWrapper wrapper(index_path, true); + + // Add some test geometries + auto point1_wkb = create_point_wkb(1.0, 1.0); + auto point2_wkb = create_point_wkb(2.0, 2.0); + auto point3_wkb = create_point_wkb(3.0, 3.0); + + wrapper.add_geometry(point1_wkb.data(), point1_wkb.size(), 0); + wrapper.add_geometry(point2_wkb.data(), point2_wkb.size(), 1); + wrapper.add_geometry(point3_wkb.data(), point3_wkb.size(), 2); + + wrapper.finish(); + } + + // Test loading index + { + milvus::index::RTreeIndexWrapper wrapper(index_path, false); + wrapper.load(); + + // Create a query geometry (polygon that contains points 1 and 2) + auto query_polygon_wkb = create_polygon_wkb( + {{0.0, 0.0}, {2.5, 0.0}, {2.5, 2.5}, {0.0, 2.5}, {0.0, 0.0}}); + + OGRGeometry* query_geom = nullptr; + OGRGeometryFactory::createFromWkb(query_polygon_wkb.data(), + nullptr, + &query_geom, + query_polygon_wkb.size()); + + ASSERT_NE(query_geom, nullptr); + + std::vector candidates; + wrapper.query_candidates( + milvus::proto::plan::GISFunctionFilterExpr_GISOp_Intersects, + *query_geom, + candidates); + + // Should find points 1 and 2, but not point 3 + EXPECT_EQ(candidates.size(), 2); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 0) != + candidates.end()); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 1) != + candidates.end()); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 2) == + candidates.end()); + + OGRGeometryFactory::destroyGeometry(query_geom); + } +} + +TEST_F(RTreeIndexWrapperTest, TestQueryOperations) { + std::string index_path = test_dir_ + "/test_query_index"; + + // Build index with various geometries + { + milvus::index::RTreeIndexWrapper wrapper(index_path, true); + + // Add a polygon + auto polygon_wkb = create_polygon_wkb( + {{0.0, 0.0}, {10.0, 0.0}, {10.0, 10.0}, {0.0, 10.0}, {0.0, 0.0}}); + wrapper.add_geometry(polygon_wkb.data(), polygon_wkb.size(), 0); + + // Add some points + auto point1_wkb = create_point_wkb(5.0, 5.0); // Inside polygon + auto point2_wkb = create_point_wkb(15.0, 15.0); // Outside polygon + auto point3_wkb = create_point_wkb(1.0, 1.0); // Inside polygon + + wrapper.add_geometry(point1_wkb.data(), point1_wkb.size(), 1); + wrapper.add_geometry(point2_wkb.data(), point2_wkb.size(), 2); + wrapper.add_geometry(point3_wkb.data(), point3_wkb.size(), 3); + + wrapper.finish(); + } + + // Test queries + { + milvus::index::RTreeIndexWrapper wrapper(index_path, false); + wrapper.load(); + + // Query with a small polygon that intersects with the large polygon + auto query_polygon_wkb = create_polygon_wkb( + {{4.0, 4.0}, {6.0, 4.0}, {6.0, 6.0}, {4.0, 6.0}, {4.0, 4.0}}); + + OGRGeometry* query_geom = nullptr; + OGRGeometryFactory::createFromWkb(query_polygon_wkb.data(), + nullptr, + &query_geom, + query_polygon_wkb.size()); + + ASSERT_NE(query_geom, nullptr); + + std::vector candidates; + wrapper.query_candidates( + milvus::proto::plan::GISFunctionFilterExpr_GISOp_Intersects, + *query_geom, + candidates); + + // Should find the large polygon and point1, but not point2 or point3 + EXPECT_EQ(candidates.size(), 2); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 0) != + candidates.end()); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 1) != + candidates.end()); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 2) == + candidates.end()); + EXPECT_TRUE(std::find(candidates.begin(), candidates.end(), 3) == + candidates.end()); + + OGRGeometryFactory::destroyGeometry(query_geom); + } +} + +TEST_F(RTreeIndexWrapperTest, TestInvalidWKB) { + std::string index_path = test_dir_ + "/test_invalid_wkb"; + + milvus::index::RTreeIndexWrapper wrapper(index_path, true); + + // Test with invalid WKB data + std::vector invalid_wkb = {0x01, 0x02, 0x03, 0x04}; // Invalid WKB + + // This should not crash and should handle the error gracefully + wrapper.add_geometry(invalid_wkb.data(), invalid_wkb.size(), 0); + + wrapper.finish(); +} \ No newline at end of file diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index e6722ea091a..27f7e965ab0 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -344,7 +344,8 @@ GenerateRandomSparseFloatVector(size_t rows, return tensor; } -inline OGRGeometry* makeGeometryValid(OGRGeometry* geometry) { +inline OGRGeometry* +makeGeometryValid(OGRGeometry* geometry) { if (!geometry || geometry->IsValid()) return geometry; diff --git a/internal/proxy/task_index.go b/internal/proxy/task_index.go index 69508081572..63e94e3fd2a 100644 --- a/internal/proxy/task_index.go +++ b/internal/proxy/task_index.go @@ -242,6 +242,8 @@ func (cit *createIndexTask) parseIndexParams(ctx context.Context) error { return getPrimitiveIndexType(cit.fieldSchema.ElementType), nil } else if typeutil.IsJSONType(dataType) { return Params.AutoIndexConfig.ScalarJSONIndexType.GetValue(), nil + } else if typeutil.IsGeometryType(dataType) { + return Params.AutoIndexConfig.ScalarGeometryIndexType.GetValue(), nil } return "", fmt.Errorf("create auto index on type:%s is not supported", dataType.String()) }() @@ -486,6 +488,7 @@ func checkTrain(ctx context.Context, field *schemapb.FieldSchema, indexParams ma indexParams[common.BitmapCardinalityLimitKey] = paramtable.Get().AutoIndexConfig.BitmapCardinalityLimit.GetValue() } } + checker, err := indexparamcheck.GetIndexCheckerMgrInstance().GetChecker(indexType) if err != nil { log.Ctx(ctx).Warn("Failed to get index checker", zap.String(common.IndexTypeKey, indexType)) diff --git a/internal/util/indexparamcheck/conf_adapter_mgr.go b/internal/util/indexparamcheck/conf_adapter_mgr.go index a746f423cea..d0f4e8a4877 100644 --- a/internal/util/indexparamcheck/conf_adapter_mgr.go +++ b/internal/util/indexparamcheck/conf_adapter_mgr.go @@ -56,6 +56,7 @@ func (mgr *indexCheckerMgrImpl) registerIndexChecker() { mgr.checkers[IndexTrie] = newTRIEChecker() mgr.checkers[IndexBitmap] = newBITMAPChecker() mgr.checkers[IndexHybrid] = newHYBRIDChecker() + mgr.checkers[IndexRTREE] = newRTREEChecker() mgr.checkers["marisa-trie"] = newTRIEChecker() mgr.checkers[AutoIndex] = newAUTOINDEXChecker() } diff --git a/internal/util/indexparamcheck/index_type.go b/internal/util/indexparamcheck/index_type.go index 45bdbdc747d..92fcf5256c3 100644 --- a/internal/util/indexparamcheck/index_type.go +++ b/internal/util/indexparamcheck/index_type.go @@ -33,6 +33,7 @@ const ( IndexBitmap IndexType = "BITMAP" IndexHybrid IndexType = "HYBRID" // BITMAP + INVERTED IndexINVERTED IndexType = "INVERTED" + IndexRTREE IndexType = "RTREE" AutoIndex IndexType = "AUTOINDEX" ) diff --git a/internal/util/indexparamcheck/rtree_checker.go b/internal/util/indexparamcheck/rtree_checker.go new file mode 100644 index 00000000000..d7144ad1885 --- /dev/null +++ b/internal/util/indexparamcheck/rtree_checker.go @@ -0,0 +1,49 @@ +// 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. + +package indexparamcheck + +import ( + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/pkg/v2/util/typeutil" +) + +// RTREEChecker checks if a RTREE index can be built. +type RTREEChecker struct { + scalarIndexChecker +} + +func (c *RTREEChecker) CheckTrain(dataType schemapb.DataType, params map[string]string) error { + if !typeutil.IsGeometryType(dataType) { + return fmt.Errorf("RTREE index can only be built on geometry field") + } + + return c.scalarIndexChecker.CheckTrain(dataType, params) +} + +func (c *RTREEChecker) CheckValidDataType(indexType IndexType, field *schemapb.FieldSchema) error { + dType := field.GetDataType() + if !typeutil.IsGeometryType(dType) { + return fmt.Errorf("RTREE index can only be built on geometry field, got %s", dType.String()) + } + return nil +} + +func newRTREEChecker() *RTREEChecker { + return &RTREEChecker{} +} diff --git a/internal/util/indexparamcheck/rtree_checker_test.go b/internal/util/indexparamcheck/rtree_checker_test.go new file mode 100644 index 00000000000..fb37ebcd602 --- /dev/null +++ b/internal/util/indexparamcheck/rtree_checker_test.go @@ -0,0 +1,162 @@ +// 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. + +package indexparamcheck + +import ( + "fmt" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" +) + +func TestRTREEChecker(t *testing.T) { + c := newRTREEChecker() + + t.Run("valid data type", func(t *testing.T) { + field := &schemapb.FieldSchema{ + DataType: schemapb.DataType_Geometry, + } + err := c.CheckValidDataType(IndexRTREE, field) + assert.NoError(t, err) + }) + + t.Run("invalid data type", func(t *testing.T) { + field := &schemapb.FieldSchema{ + DataType: schemapb.DataType_VarChar, + } + err := c.CheckValidDataType(IndexRTREE, field) + assert.Error(t, err) + }) + + t.Run("valid parameters with defaults", func(t *testing.T) { + params := make(map[string]string) + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.NoError(t, err) + + // Check that defaults are applied + assert.Equal(t, fmt.Sprintf("%f", DefaultRTreeFillFactor), params[RTreeFillFactor]) + assert.Equal(t, strconv.Itoa(DefaultRTreeIndexCapacity), params[RTreeIndexCapacity]) + assert.Equal(t, strconv.Itoa(DefaultRTreeLeafCapacity), params[RTreeLeafCapacity]) + assert.Equal(t, strconv.Itoa(DefaultRTreeDim), params[RTreeDim]) + assert.Equal(t, DefaultRTreeRV, params[RTreeRV]) + }) + + t.Run("valid custom parameters", func(t *testing.T) { + params := map[string]string{ + RTreeFillFactor: "0.7", + RTreeIndexCapacity: "150", + RTreeLeafCapacity: "150", + RTreeDim: "2", + RTreeRV: "RV_LINEAR", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.NoError(t, err) + }) + + t.Run("invalid fillFactor - too low", func(t *testing.T) { + params := map[string]string{ + RTreeFillFactor: "0.05", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid fillFactor - too high", func(t *testing.T) { + params := map[string]string{ + RTreeFillFactor: "1.5", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid indexCapacity - too low", func(t *testing.T) { + params := map[string]string{ + RTreeIndexCapacity: "1", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid indexCapacity - too high", func(t *testing.T) { + params := map[string]string{ + RTreeIndexCapacity: "1500", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid leafCapacity - too low", func(t *testing.T) { + params := map[string]string{ + RTreeLeafCapacity: "1", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid leafCapacity - too high", func(t *testing.T) { + params := map[string]string{ + RTreeLeafCapacity: "1500", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid dimension - too low", func(t *testing.T) { + params := map[string]string{ + RTreeDim: "1", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid dimension - too high", func(t *testing.T) { + params := map[string]string{ + RTreeDim: "4", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("invalid rv value", func(t *testing.T) { + params := map[string]string{ + RTreeRV: "INVALID_RV", + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.Error(t, err) + }) + + t.Run("valid rv values", func(t *testing.T) { + validRVs := []string{"RV_LINEAR", "RV_QUADRATIC", "RV_RSTAR"} + for _, rv := range validRVs { + params := map[string]string{ + RTreeRV: rv, + } + err := c.CheckTrain(schemapb.DataType_Geometry, params) + assert.NoError(t, err) + } + }) + + t.Run("non-geometry data type", func(t *testing.T) { + params := make(map[string]string) + err := c.CheckTrain(schemapb.DataType_VarChar, params) + assert.Error(t, err) + assert.Contains(t, err.Error(), "RTREE index can only be built on geometry field") + }) +} diff --git a/internal/util/indexparamcheck/utils.go b/internal/util/indexparamcheck/utils.go index 2717436e51a..93f45081152 100644 --- a/internal/util/indexparamcheck/utils.go +++ b/internal/util/indexparamcheck/utils.go @@ -47,6 +47,21 @@ func CheckIntByRange(params map[string]string, key string, min, max int) bool { return value >= min && value <= max } +// CheckFloatByRange checks if the float value of key in params is in range [min, max]. +func CheckFloatByRange(params map[string]string, key string, min, max float64) bool { + valueStr, ok := params[key] + if !ok { + return false + } + + value, err := strconv.ParseFloat(valueStr, 64) + if err != nil { + return false + } + + return value >= min && value <= max +} + // CheckStrByValues check whether the data corresponding to the key appears in the string slice of container. // Return false if: // 1. the key does not exist, or diff --git a/pkg/util/paramtable/autoindex_param.go b/pkg/util/paramtable/autoindex_param.go index 971ed1552da..61c7020cb81 100644 --- a/pkg/util/paramtable/autoindex_param.go +++ b/pkg/util/paramtable/autoindex_param.go @@ -46,14 +46,15 @@ type AutoIndexConfig struct { AutoIndexSearchConfig ParamItem `refreshable:"true"` AutoIndexTuningConfig ParamGroup `refreshable:"true"` - ScalarAutoIndexEnable ParamItem `refreshable:"true"` - ScalarAutoIndexParams ParamItem `refreshable:"true"` - ScalarNumericIndexType ParamItem `refreshable:"true"` - ScalarIntIndexType ParamItem `refreshable:"true"` - ScalarVarcharIndexType ParamItem `refreshable:"true"` - ScalarBoolIndexType ParamItem `refreshable:"true"` - ScalarFloatIndexType ParamItem `refreshable:"true"` - ScalarJSONIndexType ParamItem `refreshable:"true"` + ScalarAutoIndexEnable ParamItem `refreshable:"true"` + ScalarAutoIndexParams ParamItem `refreshable:"true"` + ScalarNumericIndexType ParamItem `refreshable:"true"` + ScalarIntIndexType ParamItem `refreshable:"true"` + ScalarVarcharIndexType ParamItem `refreshable:"true"` + ScalarBoolIndexType ParamItem `refreshable:"true"` + ScalarFloatIndexType ParamItem `refreshable:"true"` + ScalarJSONIndexType ParamItem `refreshable:"true"` + ScalarGeometryIndexType ParamItem `refreshable:"true"` BitmapCardinalityLimit ParamItem `refreshable:"true"` } @@ -167,7 +168,7 @@ func (p *AutoIndexConfig) init(base *BaseTable) { p.ScalarAutoIndexParams = ParamItem{ Key: "scalarAutoIndex.params.build", Version: "2.4.0", - DefaultValue: `{"int": "HYBRID","varchar": "HYBRID","bool": "BITMAP", "float": "INVERTED", "json": "INVERTED"}`, + DefaultValue: `{"int": "HYBRID","varchar": "HYBRID","bool": "BITMAP", "float": "INVERTED", "json": "INVERTED", "geometry": "RTREE"}`, } p.ScalarAutoIndexParams.Init(base.mgr) @@ -220,6 +221,18 @@ func (p *AutoIndexConfig) init(base *BaseTable) { } p.ScalarJSONIndexType.Init(base.mgr) + p.ScalarGeometryIndexType = ParamItem{ + Version: "2.5.16", + Formatter: func(v string) string { + m := p.ScalarAutoIndexParams.GetAsJSONMap() + if m == nil { + return "" + } + return m["geometry"] + }, + } + p.ScalarGeometryIndexType.Init(base.mgr) + p.BitmapCardinalityLimit = ParamItem{ Key: "scalarAutoIndex.params.bitmapCardinalityLimit", Version: "2.5.0",