From cdbaeb3982c297b4f1109adc659598c9d2331b9b Mon Sep 17 00:00:00 2001 From: Yinwei Li Date: Wed, 30 Jul 2025 11:26:47 +0000 Subject: [PATCH 1/2] Add geometry support. Support geometry in milvus.Now can insert,query and search with gis functions. This commit includes file changing with related to 3rdparty dependencies,cpp core and go. Here is also an example file which shows how to use the geo type and gis functions. Signed-off-by: Yinwei Li --- client/index/common.go | 1 + client/index/rtree.go | 148 +++ client/index/rtree_test.go | 119 +++ internal/core/conanfile.py | 3 +- internal/core/src/common/Geometry.h | 15 + .../exec/expression/GISFunctionFilterExpr.cpp | 220 ++++- .../exec/expression/GISFunctionFilterExpr.h | 15 +- internal/core/src/index/IndexFactory.cpp | 14 + internal/core/src/index/IndexFactory.h | 6 + internal/core/src/index/Meta.h | 6 + internal/core/src/index/RTreeIndex.cpp | 542 +++++++++++ internal/core/src/index/RTreeIndex.h | 172 ++++ internal/core/src/index/RTreeIndexWrapper.cpp | 477 ++++++++++ internal/core/src/index/RTreeIndexWrapper.h | 152 +++ internal/core/src/index/ScalarIndex.h | 3 + internal/core/src/index/Utils.cpp | 50 + internal/core/src/index/Utils.h | 9 + internal/core/src/indexbuilder/IndexFactory.h | 1 + internal/core/unittest/CMakeLists.txt | 2 + internal/core/unittest/test_expr.cpp | 2 +- internal/core/unittest/test_rtree_index.cpp | 866 ++++++++++++++++++ .../unittest/test_rtree_index_wrapper.cpp | 232 +++++ internal/core/unittest/test_utils/DataGen.h | 3 +- internal/proxy/task_index.go | 16 + .../util/indexparamcheck/conf_adapter_mgr.go | 1 + internal/util/indexparamcheck/constraints.go | 26 + internal/util/indexparamcheck/index_type.go | 1 + .../util/indexparamcheck/rtree_checker.go | 86 ++ .../indexparamcheck/rtree_checker_test.go | 162 ++++ internal/util/indexparamcheck/utils.go | 15 + pkg/util/paramtable/autoindex_param.go | 41 +- 31 files changed, 3385 insertions(+), 21 deletions(-) create mode 100644 client/index/rtree.go create mode 100644 client/index/rtree_test.go create mode 100644 internal/core/src/index/RTreeIndex.cpp create mode 100644 internal/core/src/index/RTreeIndex.h create mode 100644 internal/core/src/index/RTreeIndexWrapper.cpp create mode 100644 internal/core/src/index/RTreeIndexWrapper.h create mode 100644 internal/core/unittest/test_rtree_index.cpp create mode 100644 internal/core/unittest/test_rtree_index_wrapper.cpp create mode 100644 internal/util/indexparamcheck/rtree_checker.go create mode 100644 internal/util/indexparamcheck/rtree_checker_test.go 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..2fa6d58a87d --- /dev/null +++ b/client/index/rtree.go @@ -0,0 +1,148 @@ +// 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 ( + "strconv" +) + +// RTree index parameter keys +const ( + RTreeFillFactorKey = "fillFactor" + RTreeIndexCapacityKey = "indexCapacity" + RTreeLeafCapacityKey = "leafCapacity" + RTreeDimKey = "dim" + RTreeRVKey = "rv" +) + +// RTree index parameter defaults +const ( + DefaultRTreeFillFactor = 0.8 + DefaultRTreeIndexCapacity = 100 + DefaultRTreeLeafCapacity = 100 + DefaultRTreeDim = 2 + DefaultRTreeRV = "RSTAR" +) + +var _ Index = rtreeIndex{} + +// rtreeIndex represents an RTree index for geometry fields +type rtreeIndex struct { + baseIndex + fillFactor float64 + indexCapacity int + leafCapacity int + dim int + rv string +} + +func (idx rtreeIndex) Params() map[string]string { + params := map[string]string{ + IndexTypeKey: string(RTREE), + RTreeFillFactorKey: strconv.FormatFloat(idx.fillFactor, 'f', -1, 64), + RTreeIndexCapacityKey: strconv.Itoa(idx.indexCapacity), + RTreeLeafCapacityKey: strconv.Itoa(idx.leafCapacity), + RTreeDimKey: strconv.Itoa(idx.dim), + RTreeRVKey: idx.rv, + } + return params +} + +// NewRTreeIndex creates a new RTree index with default parameters +func NewRTreeIndex() Index { + return rtreeIndex{ + baseIndex: baseIndex{ + indexType: RTREE, + }, + fillFactor: DefaultRTreeFillFactor, + indexCapacity: DefaultRTreeIndexCapacity, + leafCapacity: DefaultRTreeLeafCapacity, + dim: DefaultRTreeDim, + rv: DefaultRTreeRV, + } +} + +// NewRTreeIndexWithParams creates a new RTree index with custom parameters +func NewRTreeIndexWithParams(fillFactor float64, indexCapacity, leafCapacity, dim int, rv string) Index { + return rtreeIndex{ + baseIndex: baseIndex{ + indexType: RTREE, + }, + fillFactor: fillFactor, + indexCapacity: indexCapacity, + leafCapacity: leafCapacity, + dim: dim, + rv: rv, + } +} + +// 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, + }, + fillFactor: DefaultRTreeFillFactor, + indexCapacity: DefaultRTreeIndexCapacity, + leafCapacity: DefaultRTreeLeafCapacity, + dim: DefaultRTreeDim, + rv: DefaultRTreeRV, + }, + } +} + +// WithFillFactor sets the fill factor for the RTree index +func (b *RTreeIndexBuilder) WithFillFactor(fillFactor float64) *RTreeIndexBuilder { + b.index.fillFactor = fillFactor + return b +} + +// WithIndexCapacity sets the index capacity for the RTree index +func (b *RTreeIndexBuilder) WithIndexCapacity(capacity int) *RTreeIndexBuilder { + b.index.indexCapacity = capacity + return b +} + +// WithLeafCapacity sets the leaf capacity for the RTree index +func (b *RTreeIndexBuilder) WithLeafCapacity(capacity int) *RTreeIndexBuilder { + b.index.leafCapacity = capacity + return b +} + +// WithDimension sets the dimension for the RTree index +func (b *RTreeIndexBuilder) WithDimension(dim int) *RTreeIndexBuilder { + b.index.dim = dim + return b +} + +// WithRVType sets the RV type for the RTree index +// Valid values: "LINEAR", "QUADRATIC", "RSTAR" +func (b *RTreeIndexBuilder) WithRVType(rv string) *RTreeIndexBuilder { + b.index.rv = rv + return b +} + +// 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..736601f6b00 --- /dev/null +++ b/client/index/rtree_test.go @@ -0,0 +1,119 @@ +// 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 ( + "strconv" + "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]) + s.Equal(strconv.FormatFloat(DefaultRTreeFillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) + s.Equal(strconv.Itoa(DefaultRTreeIndexCapacity), params[RTreeIndexCapacityKey]) + s.Equal(strconv.Itoa(DefaultRTreeLeafCapacity), params[RTreeLeafCapacityKey]) + s.Equal(strconv.Itoa(DefaultRTreeDim), params[RTreeDimKey]) + s.Equal(DefaultRTreeRV, params[RTreeRVKey]) +} + +func (s *RTreeIndexSuite) TestNewRTreeIndexWithParams() { + fillFactor := 0.7 + indexCapacity := 150 + leafCapacity := 150 + dim := 3 + rv := "RV_LINEAR" + + idx := NewRTreeIndexWithParams(fillFactor, indexCapacity, leafCapacity, dim, rv) + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) + s.Equal(strconv.FormatFloat(fillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) + s.Equal(strconv.Itoa(indexCapacity), params[RTreeIndexCapacityKey]) + s.Equal(strconv.Itoa(leafCapacity), params[RTreeLeafCapacityKey]) + s.Equal(strconv.Itoa(dim), params[RTreeDimKey]) + s.Equal(rv, params[RTreeRVKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilder() { + idx := NewRTreeIndexBuilder(). + WithFillFactor(0.6). + WithIndexCapacity(200). + WithLeafCapacity(200). + WithDimension(2). + WithRVType("RV_QUADRATIC"). + Build() + + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) + s.Equal("0.6", params[RTreeFillFactorKey]) + s.Equal("200", params[RTreeIndexCapacityKey]) + s.Equal("200", params[RTreeLeafCapacityKey]) + s.Equal("2", params[RTreeDimKey]) + s.Equal("RV_QUADRATIC", params[RTreeRVKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilderDefaults() { + idx := NewRTreeIndexBuilder().Build() + s.Equal(RTREE, idx.IndexType()) + + params := idx.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) + s.Equal(strconv.FormatFloat(DefaultRTreeFillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) + s.Equal(strconv.Itoa(DefaultRTreeIndexCapacity), params[RTreeIndexCapacityKey]) + s.Equal(strconv.Itoa(DefaultRTreeLeafCapacity), params[RTreeLeafCapacityKey]) + s.Equal(strconv.Itoa(DefaultRTreeDim), params[RTreeDimKey]) + s.Equal(DefaultRTreeRV, params[RTreeRVKey]) +} + +func (s *RTreeIndexSuite) TestRTreeIndexBuilderChaining() { + builder := NewRTreeIndexBuilder() + + // Test method chaining + result := builder. + WithFillFactor(0.9). + WithIndexCapacity(50). + WithLeafCapacity(25). + WithDimension(3). + WithRVType("RV_RSTAR") + + s.Equal(builder, result) // Should return the same builder instance + + idx := result.Build() + params := idx.Params() + s.Equal("0.9", params[RTreeFillFactorKey]) + s.Equal("50", params[RTreeIndexCapacityKey]) + s.Equal("25", params[RTreeLeafCapacityKey]) + s.Equal("3", params[RTreeDimKey]) + s.Equal("RV_RSTAR", params[RTreeRVKey]) +} + +func TestRTreeIndex(t *testing.T) { + suite.Run(t, new(RTreeIndexSuite)) +} diff --git a/internal/core/conanfile.py b/internal/core/conanfile.py index 5cb7b9d1b7c..c3db43b9ea2 100644 --- a/internal/core/conanfile.py +++ b/internal/core/conanfile.py @@ -54,7 +54,8 @@ class MilvusConan(ConanFile): "libtiff/4.6.0#32ca1d04c9f024637d49c0c2882cfdbe", "libgeotiff/1.7.1#0375633ef1116fc067b3773be7fd902f", "geos/3.12.0#b76c27884c1fa4ee8c9e486337b7dc4e", - "gdal/3.5.3#61a42c933d3440a449cac89fd0866621" + "gdal/3.5.3#61a42c933d3440a449cac89fd0866621", + "libspatialindex/2.1.0#866b4d23930c42221f0f28547ed2b3d5" ) 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..bc36078935e 100644 --- a/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp +++ b/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp @@ -49,8 +49,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 +142,219 @@ 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); + + { + LOG_INFO("LiYinwei:Query segment id {} start", + segment_->get_segment_id()); + LOG_INFO("LiYinwei:Query op {}", + ds->Get( + milvus::index::OPERATOR_TYPE)); + auto tmp = idx_ptr->Query(ds); + LOG_INFO("LiYinwei:Query segment id {} end", + segment_->get_segment_id()); + 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; + auto num_chunk_data = segment_->num_chunk_data(field_id_); + for (size_t i = current_index_chunk_; i < num_chunk_data; ++i) { + // 1) Build and cache refined bitmap for this chunk (coarse + exact) + if (cached_index_chunk_id_ != static_cast(i)) { + // Reuse segment-level coarse bitmap directly (same for all chunks) + auto& coarse = this->coarse_global_; + auto& chunk_valid = this->coarse_valid_global_; + + // Exact refinement + TargetBitmap refined(coarse.size()); + const bool is_sealed = segment_->type() == SegmentType::Sealed; + + if (is_sealed) { + auto [views, valid_vec] = + segment_->chunk_view(field_id_, i); + + // Align global coarse bitmap positions with per-chunk local views + const auto start_pos = + segment_->num_rows_until_chunk(field_id_, i); + const auto chunk_rows = views.size(); + 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; + if (!valid_vec.empty() && !valid_vec[local]) + continue; + + const auto& wkb_view = views[local]; + Geometry left(wkb_view.data(), wkb_view.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; + } + + // CRITICAL FIX: Ensure the returned ColumnVector exactly matches the real_batch_size + // This handles the case where the loop might have accumulated slightly more + // due to chunking/batching logic not perfectly aligning with `real_batch_size`. + if (batch_result.size() > real_batch_size) { + LOG_WARN( + "EvalForIndexSegment: Truncating batch_result from {} to " + "{} to match real_batch_size.", + batch_result.size(), + real_batch_size); + batch_result.resize(real_batch_size); + batch_valid.resize(real_batch_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..12a85332490 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"; @@ -93,4 +94,9 @@ constexpr const char* DISK_ANN_PREPARE_USE_BFS_CACHE = "use_bfs_cache"; // DiskAnn query params constexpr const char* DISK_ANN_QUERY_LIST = "search_list"; constexpr const char* DISK_ANN_QUERY_BEAMWIDTH = "beamwidth"; + +constexpr const char* R_TREE_VARIANT_KEY = "rv"; +constexpr const char* FILL_FACTOR_KEY = "fillFactor"; +constexpr const char* INDEX_CAPACITY_KEY = "indexCapacity"; +constexpr const char* LEAF_CAPACITY_KEY = "leafCapacity"; } // namespace milvus::index diff --git a/internal/core/src/index/RTreeIndex.cpp b/internal/core/src/index/RTreeIndex.cpp new file mode 100644 index 00000000000..fd53532c58e --- /dev/null +++ b/internal/core/src/index/RTreeIndex.cpp @@ -0,0 +1,542 @@ +// 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 +#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, ".dat")) { + base_path = p.substr(0, p.size() - 4); + break; + } + if (ends_with(p, ".idx")) { + 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(); + auto fill_factor = GetFillFactorFromConfig(config); + auto index_cap = GetIndexCapacityFromConfig(config); + auto leaf_cap = GetLeafCapacityFromConfig(config); + auto variant_str = + GetValueFromConfig(config, R_TREE_VARIANT_KEY) + .value_or("RSTAR"); + wrapper_->set_fill_factor(fill_factor); + wrapper_->set_index_capacity(index_cap); + wrapper_->set_leaf_capacity(leaf_cap); + wrapper_->set_rtree_variant(variant_str); + + // 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..d47f8ea5182 --- /dev/null +++ b/internal/core/src/index/RTreeIndex.h @@ -0,0 +1,172 @@ +// 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/RTreeIndexWrapper.cpp b/internal/core/src/index/RTreeIndexWrapper.cpp new file mode 100644 index 00000000000..3cf499126f0 --- /dev/null +++ b/internal/core/src/index/RTreeIndexWrapper.cpp @@ -0,0 +1,477 @@ +// 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 "common/EasyAssert.h" +#include "log/Log.h" +#include "pb/plan.pb.h" +#include +#include +#include +#include "common/FieldDataInterface.h" + +namespace milvus::index { + +// Custom visitor for collecting query results +class GeometryVisitor : public SpatialIndex::IVisitor { + public: + explicit GeometryVisitor(std::vector& results) + : results_(results) { + } + + virtual ~GeometryVisitor() = default; + + void + visitNode(const SpatialIndex::INode& n) override { + // Not needed for our use case + } + + void + visitData(const SpatialIndex::IData& d) override { + // Store the identifier (row offset) in results + results_.push_back(static_cast(d.getIdentifier())); + } + + void + visitData(std::vector& v) override { + for (const auto* data : v) { + results_.push_back(static_cast(data->getIdentifier())); + } + } + + private: + std::vector& results_; +}; + +RTreeIndexWrapper::RTreeIndexWrapper(std::string& path, bool is_build_mode) + : index_path_(path), is_build_mode_(is_build_mode) { + if (is_build_mode_) { + // Create directory if it doesn't exist + std::filesystem::path dir_path = + std::filesystem::path(path).parent_path(); + if (!dir_path.empty()) { + std::filesystem::create_directories(dir_path); + } + + // Create disk storage manager for building + storage_manager_ = std::shared_ptr( + SpatialIndex::StorageManager::createNewDiskStorageManager(path, + 4096)); + } +} + +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"); + // Lazily create the R-Tree for dynamic insertion if not present yet + if (rtree_ == nullptr) { + SpatialIndex::id_type index_id; + rtree_ = std::shared_ptr( + SpatialIndex::RTree::createNewRTree(*storage_manager_, + fill_factor_, + index_capacity_, + leaf_capacity_, + dimension_, + rtree_variant_, + index_id)); + index_id_ = index_id; + LOG_WARN("create rtree index for dynamic insertion"); + } + + // 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 region for the bounding box + double low[2] = {minX, minY}; + double high[2] = {maxX, maxY}; + SpatialIndex::Region region(low, high, 2); + + // Insert into R-Tree with row_offset as identifier + rtree_->insertData( + 0, nullptr, region, static_cast(row_offset)); + + // Clean up + OGRGeometryFactory::destroyGeometry(geom); +} + +// Internal IDataStream implementation over FieldDataBase (WKB string rows) +namespace { +class BulkLoadDataStream : public SpatialIndex::IDataStream { + public: + BulkLoadDataStream( + const std::vector>& + field_datas, + bool nullable) + : field_datas_(field_datas), nullable_param_(nullable) { + // Compute a cheap upper bound for stream size: sum of row counts + total_rows_ = 0; + for (const auto& fd : field_datas_) { + total_rows_ += static_cast(fd->get_num_rows()); + } + rewind(); + } + + ~BulkLoadDataStream() override = default; + + bool + hasNext() override { + return absolute_offset_ < static_cast(total_rows_); + } + + uint32_t + size() override { + // Return upper bound; actual yielded items may be fewer due to + // null rows or invalid WKB filtered in getNext(). + return static_cast(total_rows_); + } + + void + rewind() override { + batch_index_ = 0; + row_in_batch_ = 0; + absolute_offset_ = 0; + } + + SpatialIndex::IData* + getNext() override { + while (batch_index_ < field_datas_.size()) { + const auto& fd = field_datas_[batch_index_]; + auto n = fd->get_num_rows(); + if (row_in_batch_ >= n) { + ++batch_index_; + row_in_batch_ = 0; + continue; + } + + int64_t current_row_in_batch = row_in_batch_; + int64_t current_abs = absolute_offset_; + // advance offsets for next call regardless of validity + ++row_in_batch_; + ++absolute_offset_; + + const bool is_nullable_effective = + nullable_param_ || fd->IsNullable(); + if (is_nullable_effective && !fd->is_valid(current_row_in_batch)) { + continue; + } + + const auto* wkb_str = static_cast( + fd->RawValue(current_row_in_batch)); + if (wkb_str == nullptr || wkb_str->empty()) { + continue; + } + + // Parse WKB using OGR to get envelope + OGRGeometry* geom = nullptr; + OGRErr err = OGRGeometryFactory::createFromWkb( + reinterpret_cast(wkb_str->data()), + nullptr, + &geom, + wkb_str->size()); + if (err != OGRERR_NONE || geom == nullptr) { + LOG_WARN( + "BulkLoadDataStream: failed to parse WKB at abs {} (batch " + "{}, row {})", + current_abs, + batch_index_, + current_row_in_batch); + continue; + } + + OGREnvelope env; + geom->getEnvelope(&env); + OGRGeometryFactory::destroyGeometry(geom); + + double low[2] = {env.MinX, env.MinY}; + double high[2] = {env.MaxX, env.MaxY}; + SpatialIndex::Region region(low, high, 2); + + return new SpatialIndex::RTree::Data( + 0, + nullptr, + region, + static_cast(current_abs)); + } + return nullptr; + } + + private: + const std::vector>& field_datas_; + bool nullable_param_ = false; + size_t total_rows_ = 0; + size_t batch_index_ = 0; + int64_t row_in_batch_ = 0; + int64_t absolute_offset_ = 0; +}; +} // anonymous namespace + +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"); + AssertInfo(storage_manager_ != nullptr, "Storage manager is null"); + AssertInfo(rtree_ == nullptr, + "R-Tree already initialized; bulk load requires a fresh tree"); + + BulkLoadDataStream stream(field_datas, nullable); + SpatialIndex::id_type index_id; + try { + rtree_ = std::shared_ptr( + SpatialIndex::RTree::createAndBulkLoadNewRTree( + SpatialIndex::RTree::BLM_STR, + stream, + *storage_manager_, + fill_factor_, + index_capacity_, + leaf_capacity_, + dimension_, + rtree_variant_, + index_id)); + } catch (const std::exception& e) { + LOG_ERROR("Failed to bulk load R-Tree: {}", e.what()); + } + + index_id_ = index_id; + LOG_INFO("R-Tree bulk load completed with {} entries", + rtree_ ? "some" : "none"); +} + +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). + if (finished_) { + LOG_DEBUG("RTreeIndexWrapper::finish() called more than once, skip."); + return; + } + + AssertInfo(is_build_mode_, "Cannot finish in load mode"); + + // If rtree_ is already reset, we have nothing left to do. Mark finished + // and return. + if (rtree_ == nullptr) { + LOG_DEBUG( + "RTreeIndexWrapper::finish() called with null rtree_, likely " + "already finished."); + finished_ = true; + return; + } + + // Explicitly flush the index header & buffers to disk to guarantee + // consistency before releasing resources. + rtree_->flush(); + + // NOTE: rtree_ internally holds a pointer to the storage manager. We must + // make sure rtree_ is destroyed BEFORE the storage manager. + + // 1. Release rtree_ first so its destructor can safely write the header + // using a still-valid storage_manager_. + rtree_.reset(); + + // 2. Now it is safe to release the storage manager. + storage_manager_.reset(); + + // 3. Write meta file with index parameters for reliable loading. + try { + nlohmann::json meta; + meta["index_id"] = index_id_; + meta["variant"] = static_cast(rtree_variant_); + meta["fill_factor"] = fill_factor_; + meta["index_capacity"] = index_capacity_; + meta["leaf_capacity"] = leaf_capacity_; + meta["dimension"] = dimension_; + + 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 meta json: {}", e.what()); + } + + finished_ = true; + + LOG_INFO("R-Tree index 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"); + + try { + // Load storage manager + storage_manager_ = std::shared_ptr( + SpatialIndex::StorageManager::loadDiskStorageManager(index_path_)); + + // Determine index id from meta json if available + SpatialIndex::id_type idx_id_to_load = 0; + try { + std::ifstream ifs(index_path_ + ".meta.json"); + if (ifs.good()) { + auto meta = nlohmann::json::parse(ifs); + if (meta.contains("index_id")) { + idx_id_to_load = + meta["index_id"].get(); + } + } + } catch (const std::exception& e) { + LOG_WARN("Failed to read meta json, fallback to default id 0: {}", + e.what()); + } + + // Load R-Tree index with the resolved id + rtree_ = std::shared_ptr( + SpatialIndex::RTree::loadRTree(*storage_manager_, idx_id_to_load)); + + LOG_INFO("R-Tree index 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) { + // Acquire read lock to protect rtree_ access during query + folly::SharedMutexWritePriority::ReadHolder lock(rtree_mutex_); + + AssertInfo(rtree_ != nullptr, "R-Tree index not initialized"); + + 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 region + double low[2] = {minX, minY}; + double high[2] = {maxX, maxY}; + SpatialIndex::Region query_region(low, high, 2); + + // Create visitor for collecting results + GeometryVisitor visitor(candidate_offsets); + + // Perform query based on operation type + switch (op) { + default: + // For all GIS operations, we use intersection query as coarse filtering + // The exact geometric relationship will be checked in the refinement phase + rtree_->intersectsWithQuery(query_region, visitor); + break; + } + + 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 { + // Acquire read lock to protect rtree_ access during count operation + folly::SharedMutexWritePriority::ReadHolder lock(rtree_mutex_); + + if (rtree_ == nullptr) { + return 0; + } + + // For R-Tree, we need to count the number of data entries + // This is a simplified implementation - in practice, you might want to + // maintain a separate counter during building + SpatialIndex::IStatistics* stats = nullptr; + rtree_->getStatistics(&stats); + if (stats != nullptr) { + int64_t count = stats->getNumberOfData(); + delete stats; + return count; + } + return 0; +} + +void +RTreeIndexWrapper::set_rtree_variant(const std::string& variant_str) { + if (variant_str == "RSTAR") { + rtree_variant_ = SpatialIndex::RTree::RV_RSTAR; + } else if (variant_str == "QUADRATIC") { + LOG_WARN("QUADRATIC variant is not supported, using RSTAR instead"); + rtree_variant_ = SpatialIndex::RTree::RV_RSTAR; + } else if (variant_str == "LINEAR") { + rtree_variant_ = SpatialIndex::RTree::RV_LINEAR; + } else { + PanicInfo(ErrorCode::UnexpectedError, + fmt::format("Invalid R-Tree variant: {}", variant_str)); + } +} + +void +RTreeIndexWrapper::set_fill_factor(double fill_factor) { + fill_factor_ = fill_factor; +} + +void +RTreeIndexWrapper::set_index_capacity(uint32_t index_capacity) { + index_capacity_ = index_capacity; +} + +void +RTreeIndexWrapper::set_leaf_capacity(uint32_t leaf_capacity) { + leaf_capacity_ = leaf_capacity; +} +} // 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..9a2f3a201b9 --- /dev/null +++ b/internal/core/src/index/RTreeIndexWrapper.h @@ -0,0 +1,152 @@ +// 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 "ogr_geometry.h" +#include "spatialindex/SpatialIndex.h" +#include "common/Types.h" +#include "pb/plan.pb.h" +#include + +// Forward declaration to avoid pulling heavy field data headers here +namespace milvus { +class FieldDataBase; +} + +namespace milvus::index { + +/** + * @brief Wrapper class for libspatialindex R-Tree functionality + * + * This class provides a simplified interface to libspatialindex 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 to the index + * @param wkb_data Pointer to WKB binary data + * @param len Length of WKB data + * @param row_offset Row offset (used as identifier) + */ + 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; + + void + set_rtree_variant(const std::string& variant_str); + + void + set_fill_factor(double fill_factor); + + void + set_index_capacity(uint32_t index_capacity); + + void + set_leaf_capacity(uint32_t leaf_capacity); + + 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: + std::shared_ptr storage_manager_; + std::shared_ptr rtree_; + 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; + SpatialIndex::id_type index_id_ = 0; // persisted to meta for reliable load + + // R-Tree parameters + double fill_factor_ = 0.8; + uint32_t index_capacity_ = 50; + uint32_t leaf_capacity_ = 50; + uint32_t dimension_ = 2; + SpatialIndex::RTree::RTreeVariant rtree_variant_ = + SpatialIndex::RTree::RV_RSTAR; + + // Thread safety: protects rtree_ and related operations + mutable folly::SharedMutexWritePriority rtree_mutex_; +}; + +} // 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/index/Utils.cpp b/internal/core/src/index/Utils.cpp index 0abef5fbbe6..859f8736805 100644 --- a/internal/core/src/index/Utils.cpp +++ b/internal/core/src/index/Utils.cpp @@ -251,6 +251,56 @@ GetIndexMetaFromConfig(const Config& config) { return index_meta; } +double +GetFillFactorFromConfig(const Config& config) { + auto fill_factor = GetValueFromConfig(config, FILL_FACTOR_KEY); + AssertInfo(fill_factor.has_value(), + "fill factor not exist in index config"); + try { + return (std::stod(fill_factor.value())); + } catch (const std::logic_error& e) { + auto err_message = fmt::format("invalided fill factor:{}, error:{}", + fill_factor.value(), + e.what()); + LOG_ERROR(err_message); + throw std::logic_error(err_message); + } +} + +uint32_t +GetIndexCapacityFromConfig(const Config& config) { + auto index_capacity = + GetValueFromConfig(config, INDEX_CAPACITY_KEY); + AssertInfo(index_capacity.has_value(), + "index capacity not exist in index config"); + try { + return (std::stoi(index_capacity.value())); + } catch (const std::logic_error& e) { + auto err_message = fmt::format("invalided index capacity:{}, error:{}", + index_capacity.value(), + e.what()); + LOG_ERROR(err_message); + throw std::logic_error(err_message); + } +} + +uint32_t +GetLeafCapacityFromConfig(const Config& config) { + auto leaf_capacity = + GetValueFromConfig(config, LEAF_CAPACITY_KEY); + AssertInfo(leaf_capacity.has_value(), + "leaf capacity not exist in index config"); + try { + return (std::stoi(leaf_capacity.value())); + } catch (const std::logic_error& e) { + auto err_message = fmt::format("invalided leaf capacity:{}, error:{}", + leaf_capacity.value(), + e.what()); + LOG_ERROR(err_message); + throw std::logic_error(err_message); + } +} + Config ParseConfigFromIndexParams( const std::map& index_params) { diff --git a/internal/core/src/index/Utils.h b/internal/core/src/index/Utils.h index 97067b48eab..13bd84dae6f 100644 --- a/internal/core/src/index/Utils.h +++ b/internal/core/src/index/Utils.h @@ -163,6 +163,15 @@ GetFieldDataMetaFromConfig(const Config& config); storage::IndexMeta GetIndexMetaFromConfig(const Config& config); +double +GetFillFactorFromConfig(const Config& config); + +uint32_t +GetIndexCapacityFromConfig(const Config& config); + +uint32_t +GetLeafCapacityFromConfig(const Config& config); + Config ParseConfigFromIndexParams( const std::map& index_params); 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..a0ecbcd6c54 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,20 @@ func checkTrain(ctx context.Context, field *schemapb.FieldSchema, indexParams ma indexParams[common.BitmapCardinalityLimitKey] = paramtable.Get().AutoIndexConfig.BitmapCardinalityLimit.GetValue() } } + + if indexType == indexparamcheck.IndexRTREE { + // Apply default RTree parameters if not provided + rtreeParams := paramtable.Get().AutoIndexConfig.RTreeAutoIndexParams.GetAsJSONMap() + for k, v := range rtreeParams { + if k != common.IndexTypeKey { // Don't override index_type + _, exist := indexParams[k] + if !exist { + indexParams[k] = v + } + } + } + } + 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/constraints.go b/internal/util/indexparamcheck/constraints.go index cf0863d7e19..64911880221 100644 --- a/internal/util/indexparamcheck/constraints.go +++ b/internal/util/indexparamcheck/constraints.go @@ -49,6 +49,32 @@ const ( BM25B = "bm25_b" MaxBitmapCardinalityLimit = 1000 + + // RTree Index Param + RTreeFillFactor = "fillFactor" + RTreeIndexCapacity = "indexCapacity" + RTreeLeafCapacity = "leafCapacity" + RTreeDim = "dim" + RTreeRV = "rv" + + // RTree parameter constraints + MinRTreeFillFactor = 0.1 + MaxRTreeFillFactor = 1.0 + DefaultRTreeFillFactor = 0.8 + + MinRTreeIndexCapacity = 2 + MaxRTreeIndexCapacity = 1000 + DefaultRTreeIndexCapacity = 128 + + MinRTreeLeafCapacity = 2 + MaxRTreeLeafCapacity = 1000 + DefaultRTreeLeafCapacity = 128 + + MinRTreeDim = 2 + MaxRTreeDim = 2 + DefaultRTreeDim = 2 + + DefaultRTreeRV = "RSTAR" ) var ( 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..55ea1024618 --- /dev/null +++ b/internal/util/indexparamcheck/rtree_checker.go @@ -0,0 +1,86 @@ +// 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/funcutil" + "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") + } + + // Set default values if not provided + setDefaultIfNotExist(params, RTreeFillFactor, fmt.Sprintf("%f", DefaultRTreeFillFactor)) + setDefaultIfNotExist(params, RTreeIndexCapacity, fmt.Sprintf("%d", DefaultRTreeIndexCapacity)) + setDefaultIfNotExist(params, RTreeLeafCapacity, fmt.Sprintf("%d", DefaultRTreeLeafCapacity)) + setDefaultIfNotExist(params, RTreeDim, fmt.Sprintf("%d", DefaultRTreeDim)) + setDefaultIfNotExist(params, RTreeRV, DefaultRTreeRV) + + // Validate fillFactor + if !CheckFloatByRange(params, RTreeFillFactor, MinRTreeFillFactor, MaxRTreeFillFactor) { + return errOutOfRange(params[RTreeFillFactor], MinRTreeFillFactor, MaxRTreeFillFactor) + } + + // Validate indexCapacity + if !CheckIntByRange(params, RTreeIndexCapacity, MinRTreeIndexCapacity, MaxRTreeIndexCapacity) { + return errOutOfRange(params[RTreeIndexCapacity], MinRTreeIndexCapacity, MaxRTreeIndexCapacity) + } + + // Validate leafCapacity + if !CheckIntByRange(params, RTreeLeafCapacity, MinRTreeLeafCapacity, MaxRTreeLeafCapacity) { + return errOutOfRange(params[RTreeLeafCapacity], MinRTreeLeafCapacity, MaxRTreeLeafCapacity) + } + + // Validate dim + if !CheckIntByRange(params, RTreeDim, MinRTreeDim, MaxRTreeDim) { + return errOutOfRange(params[RTreeDim], MinRTreeDim, MaxRTreeDim) + } + + // Validate rv + rvValue, exists := params[RTreeRV] + if exists { + validRVValues := []string{"LINEAR", "QUADRATIC", "RSTAR"} + if !funcutil.SliceContain(validRVValues, rvValue) { + return fmt.Errorf("rv value %s is not supported, supported values: %v", rvValue, validRVValues) + } + } + + 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..5bc7faa038d 100644 --- a/pkg/util/paramtable/autoindex_param.go +++ b/pkg/util/paramtable/autoindex_param.go @@ -46,15 +46,17 @@ 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"` + + RTreeAutoIndexParams ParamItem `refreshable:"true"` BitmapCardinalityLimit ParamItem `refreshable:"true"` } @@ -167,7 +169,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 +222,25 @@ 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.RTreeAutoIndexParams = ParamItem{ + Key: "scalarAutoIndex.params.rtree", + Version: "2.5.16", + DefaultValue: `{"fillFactor": 0.8, "indexCapacity": 100, "leafCapacity": 100, "dim": 2, "rv": "RSTAR", "index_type": "RTREE"}`, + } + p.RTreeAutoIndexParams.Init(base.mgr) + p.BitmapCardinalityLimit = ParamItem{ Key: "scalarAutoIndex.params.bitmapCardinalityLimit", Version: "2.5.0", From 2e266e3c25fb19be16f5d8a042940cb42fd56c2c Mon Sep 17 00:00:00 2001 From: Cai Zhang Date: Fri, 15 Aug 2025 17:55:50 +0800 Subject: [PATCH 2/2] Debug geo index Signed-off-by: Cai Zhang --- client/index/rtree.go | 82 +--- client/index/rtree_test.go | 52 +-- internal/core/CMakeLists.txt | 2 + internal/core/conanfile.py | 5 +- .../exec/expression/GISFunctionFilterExpr.cpp | 150 ++++--- internal/core/src/index/Meta.h | 5 - internal/core/src/index/RTreeIndex.cpp | 17 +- internal/core/src/index/RTreeIndex.h | 4 +- .../core/src/index/RTreeIndexSerialization.h | 147 +++++++ internal/core/src/index/RTreeIndexWrapper.cpp | 370 ++++-------------- internal/core/src/index/RTreeIndexWrapper.h | 52 ++- internal/core/src/index/Utils.cpp | 50 --- internal/core/src/index/Utils.h | 9 - internal/proxy/task_index.go | 13 - internal/util/indexparamcheck/constraints.go | 26 -- .../util/indexparamcheck/rtree_checker.go | 37 -- pkg/util/paramtable/autoindex_param.go | 8 - 17 files changed, 330 insertions(+), 699 deletions(-) create mode 100644 internal/core/src/index/RTreeIndexSerialization.h diff --git a/client/index/rtree.go b/client/index/rtree.go index 2fa6d58a87d..a01c7c4995a 100644 --- a/client/index/rtree.go +++ b/client/index/rtree.go @@ -16,48 +16,16 @@ package index -import ( - "strconv" -) - -// RTree index parameter keys -const ( - RTreeFillFactorKey = "fillFactor" - RTreeIndexCapacityKey = "indexCapacity" - RTreeLeafCapacityKey = "leafCapacity" - RTreeDimKey = "dim" - RTreeRVKey = "rv" -) - -// RTree index parameter defaults -const ( - DefaultRTreeFillFactor = 0.8 - DefaultRTreeIndexCapacity = 100 - DefaultRTreeLeafCapacity = 100 - DefaultRTreeDim = 2 - DefaultRTreeRV = "RSTAR" -) - var _ Index = rtreeIndex{} // rtreeIndex represents an RTree index for geometry fields type rtreeIndex struct { baseIndex - fillFactor float64 - indexCapacity int - leafCapacity int - dim int - rv string } func (idx rtreeIndex) Params() map[string]string { params := map[string]string{ - IndexTypeKey: string(RTREE), - RTreeFillFactorKey: strconv.FormatFloat(idx.fillFactor, 'f', -1, 64), - RTreeIndexCapacityKey: strconv.Itoa(idx.indexCapacity), - RTreeLeafCapacityKey: strconv.Itoa(idx.leafCapacity), - RTreeDimKey: strconv.Itoa(idx.dim), - RTreeRVKey: idx.rv, + IndexTypeKey: string(RTREE), } return params } @@ -68,25 +36,15 @@ func NewRTreeIndex() Index { baseIndex: baseIndex{ indexType: RTREE, }, - fillFactor: DefaultRTreeFillFactor, - indexCapacity: DefaultRTreeIndexCapacity, - leafCapacity: DefaultRTreeLeafCapacity, - dim: DefaultRTreeDim, - rv: DefaultRTreeRV, } } // NewRTreeIndexWithParams creates a new RTree index with custom parameters -func NewRTreeIndexWithParams(fillFactor float64, indexCapacity, leafCapacity, dim int, rv string) Index { +func NewRTreeIndexWithParams() Index { return rtreeIndex{ baseIndex: baseIndex{ indexType: RTREE, }, - fillFactor: fillFactor, - indexCapacity: indexCapacity, - leafCapacity: leafCapacity, - dim: dim, - rv: rv, } } @@ -102,46 +60,10 @@ func NewRTreeIndexBuilder() *RTreeIndexBuilder { baseIndex: baseIndex{ indexType: RTREE, }, - fillFactor: DefaultRTreeFillFactor, - indexCapacity: DefaultRTreeIndexCapacity, - leafCapacity: DefaultRTreeLeafCapacity, - dim: DefaultRTreeDim, - rv: DefaultRTreeRV, }, } } -// WithFillFactor sets the fill factor for the RTree index -func (b *RTreeIndexBuilder) WithFillFactor(fillFactor float64) *RTreeIndexBuilder { - b.index.fillFactor = fillFactor - return b -} - -// WithIndexCapacity sets the index capacity for the RTree index -func (b *RTreeIndexBuilder) WithIndexCapacity(capacity int) *RTreeIndexBuilder { - b.index.indexCapacity = capacity - return b -} - -// WithLeafCapacity sets the leaf capacity for the RTree index -func (b *RTreeIndexBuilder) WithLeafCapacity(capacity int) *RTreeIndexBuilder { - b.index.leafCapacity = capacity - return b -} - -// WithDimension sets the dimension for the RTree index -func (b *RTreeIndexBuilder) WithDimension(dim int) *RTreeIndexBuilder { - b.index.dim = dim - return b -} - -// WithRVType sets the RV type for the RTree index -// Valid values: "LINEAR", "QUADRATIC", "RSTAR" -func (b *RTreeIndexBuilder) WithRVType(rv string) *RTreeIndexBuilder { - b.index.rv = rv - return b -} - // 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 index 736601f6b00..7f0bcfc3881 100644 --- a/client/index/rtree_test.go +++ b/client/index/rtree_test.go @@ -17,7 +17,6 @@ package index import ( - "strconv" "testing" "github.com/stretchr/testify/suite" @@ -33,50 +32,24 @@ func (s *RTreeIndexSuite) TestNewRTreeIndex() { params := idx.Params() s.Equal(string(RTREE), params[IndexTypeKey]) - s.Equal(strconv.FormatFloat(DefaultRTreeFillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) - s.Equal(strconv.Itoa(DefaultRTreeIndexCapacity), params[RTreeIndexCapacityKey]) - s.Equal(strconv.Itoa(DefaultRTreeLeafCapacity), params[RTreeLeafCapacityKey]) - s.Equal(strconv.Itoa(DefaultRTreeDim), params[RTreeDimKey]) - s.Equal(DefaultRTreeRV, params[RTreeRVKey]) } func (s *RTreeIndexSuite) TestNewRTreeIndexWithParams() { - fillFactor := 0.7 - indexCapacity := 150 - leafCapacity := 150 - dim := 3 - rv := "RV_LINEAR" - - idx := NewRTreeIndexWithParams(fillFactor, indexCapacity, leafCapacity, dim, rv) + idx := NewRTreeIndexWithParams() s.Equal(RTREE, idx.IndexType()) params := idx.Params() s.Equal(string(RTREE), params[IndexTypeKey]) - s.Equal(strconv.FormatFloat(fillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) - s.Equal(strconv.Itoa(indexCapacity), params[RTreeIndexCapacityKey]) - s.Equal(strconv.Itoa(leafCapacity), params[RTreeLeafCapacityKey]) - s.Equal(strconv.Itoa(dim), params[RTreeDimKey]) - s.Equal(rv, params[RTreeRVKey]) } func (s *RTreeIndexSuite) TestRTreeIndexBuilder() { idx := NewRTreeIndexBuilder(). - WithFillFactor(0.6). - WithIndexCapacity(200). - WithLeafCapacity(200). - WithDimension(2). - WithRVType("RV_QUADRATIC"). Build() s.Equal(RTREE, idx.IndexType()) params := idx.Params() s.Equal(string(RTREE), params[IndexTypeKey]) - s.Equal("0.6", params[RTreeFillFactorKey]) - s.Equal("200", params[RTreeIndexCapacityKey]) - s.Equal("200", params[RTreeLeafCapacityKey]) - s.Equal("2", params[RTreeDimKey]) - s.Equal("RV_QUADRATIC", params[RTreeRVKey]) } func (s *RTreeIndexSuite) TestRTreeIndexBuilderDefaults() { @@ -85,33 +58,18 @@ func (s *RTreeIndexSuite) TestRTreeIndexBuilderDefaults() { params := idx.Params() s.Equal(string(RTREE), params[IndexTypeKey]) - s.Equal(strconv.FormatFloat(DefaultRTreeFillFactor, 'f', -1, 64), params[RTreeFillFactorKey]) - s.Equal(strconv.Itoa(DefaultRTreeIndexCapacity), params[RTreeIndexCapacityKey]) - s.Equal(strconv.Itoa(DefaultRTreeLeafCapacity), params[RTreeLeafCapacityKey]) - s.Equal(strconv.Itoa(DefaultRTreeDim), params[RTreeDimKey]) - s.Equal(DefaultRTreeRV, params[RTreeRVKey]) } func (s *RTreeIndexSuite) TestRTreeIndexBuilderChaining() { builder := NewRTreeIndexBuilder() // Test method chaining - result := builder. - WithFillFactor(0.9). - WithIndexCapacity(50). - WithLeafCapacity(25). - WithDimension(3). - WithRVType("RV_RSTAR") + result := builder.Build() - s.Equal(builder, result) // Should return the same builder instance + s.Equal(RTREE, result.IndexType()) - idx := result.Build() - params := idx.Params() - s.Equal("0.9", params[RTreeFillFactorKey]) - s.Equal("50", params[RTreeIndexCapacityKey]) - s.Equal("25", params[RTreeLeafCapacityKey]) - s.Equal("3", params[RTreeDimKey]) - s.Equal("RV_RSTAR", params[RTreeRVKey]) + params := result.Params() + s.Equal(string(RTREE), params[IndexTypeKey]) } func TestRTreeIndex(t *testing.T) { 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 c3db43b9ea2..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,9 +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", + "geos/3.12.0#0b177c90c25a8ca210578fb9e2899c37", "gdal/3.5.3#61a42c933d3440a449cac89fd0866621", - "libspatialindex/2.1.0#866b4d23930c42221f0f28547ed2b3d5" ) generators = ("cmake", "cmake_find_package") default_options = { diff --git a/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp b/internal/core/src/exec/expression/GISFunctionFilterExpr.cpp index bc36078935e..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 { @@ -167,14 +168,7 @@ PhyGISFunctionFilterExpr::EvalForIndexSegment() { auto* idx_ptr = const_cast(&idx_ref); { - LOG_INFO("LiYinwei:Query segment id {} start", - segment_->get_segment_id()); - LOG_INFO("LiYinwei:Query op {}", - ds->Get( - milvus::index::OPERATOR_TYPE)); auto tmp = idx_ptr->Query(ds); - LOG_INFO("LiYinwei:Query segment id {} end", - segment_->get_segment_id()); coarse_global_ = std::move(tmp); } { @@ -188,70 +182,88 @@ PhyGISFunctionFilterExpr::EvalForIndexSegment() { TargetBitmap batch_result; TargetBitmap batch_valid; int processed_rows = 0; - auto num_chunk_data = segment_->num_chunk_data(field_id_); - for (size_t i = current_index_chunk_; i < num_chunk_data; ++i) { + + 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 bitmap directly (same for all chunks) - auto& coarse = this->coarse_global_; - auto& chunk_valid = this->coarse_valid_global_; - + // 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) { - auto [views, valid_vec] = - segment_->chunk_view(field_id_, i); + // 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)); + } + } - // Align global coarse bitmap positions with per-chunk local views - const auto start_pos = - segment_->num_rows_until_chunk(field_id_, i); - const auto chunk_rows = views.size(); - const auto max_local = std::min( - chunk_rows, - coarse.size() > start_pos ? coarse.size() - start_pos : 0); + 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()); - for (size_t local = 0; local < max_local; ++local) { - const size_t pos = start_pos + local; - if (!coarse[pos]) - continue; - if (!valid_vec.empty() && !valid_vec[local]) - continue; + // 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(); - const auto& wkb_view = views[local]; - Geometry left(wkb_view.data(), wkb_view.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); + 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 @@ -338,20 +350,6 @@ PhyGISFunctionFilterExpr::EvalForIndexSegment() { } processed_rows += size; } - - // CRITICAL FIX: Ensure the returned ColumnVector exactly matches the real_batch_size - // This handles the case where the loop might have accumulated slightly more - // due to chunking/batching logic not perfectly aligning with `real_batch_size`. - if (batch_result.size() > real_batch_size) { - LOG_WARN( - "EvalForIndexSegment: Truncating batch_result from {} to " - "{} to match real_batch_size.", - batch_result.size(), - real_batch_size); - batch_result.resize(real_batch_size); - batch_valid.resize(real_batch_size); - } - return std::make_shared(std::move(batch_result), std::move(batch_valid)); } diff --git a/internal/core/src/index/Meta.h b/internal/core/src/index/Meta.h index 12a85332490..eacc501874c 100644 --- a/internal/core/src/index/Meta.h +++ b/internal/core/src/index/Meta.h @@ -94,9 +94,4 @@ constexpr const char* DISK_ANN_PREPARE_USE_BFS_CACHE = "use_bfs_cache"; // DiskAnn query params constexpr const char* DISK_ANN_QUERY_LIST = "search_list"; constexpr const char* DISK_ANN_QUERY_BEAMWIDTH = "beamwidth"; - -constexpr const char* R_TREE_VARIANT_KEY = "rv"; -constexpr const char* FILL_FACTOR_KEY = "fillFactor"; -constexpr const char* INDEX_CAPACITY_KEY = "indexCapacity"; -constexpr const char* LEAF_CAPACITY_KEY = "leafCapacity"; } // namespace milvus::index diff --git a/internal/core/src/index/RTreeIndex.cpp b/internal/core/src/index/RTreeIndex.cpp index fd53532c58e..6a0d31ed8b8 100644 --- a/internal/core/src/index/RTreeIndex.cpp +++ b/internal/core/src/index/RTreeIndex.cpp @@ -11,7 +11,6 @@ #include "index/RTreeIndex.h" #include -#include #include #include #include "common/Slice.h" // for INDEX_FILE_SLICE_META and Disassemble @@ -186,11 +185,7 @@ RTreeIndex::Load(milvus::tracer::TraceContext ctx, const Config& config) { // 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, ".dat")) { - base_path = p.substr(0, p.size() - 4); - break; - } - if (ends_with(p, ".idx")) { + if (ends_with(p, ".bgi")) { base_path = p.substr(0, p.size() - 4); break; } @@ -232,16 +227,6 @@ RTreeIndex::Build(const Config& config) { AssertInfo(insert_files.has_value(), "insert_files were empty for building RTree index"); InitForBuildIndex(); - auto fill_factor = GetFillFactorFromConfig(config); - auto index_cap = GetIndexCapacityFromConfig(config); - auto leaf_cap = GetLeafCapacityFromConfig(config); - auto variant_str = - GetValueFromConfig(config, R_TREE_VARIANT_KEY) - .value_or("RSTAR"); - wrapper_->set_fill_factor(fill_factor); - wrapper_->set_index_capacity(index_cap); - wrapper_->set_leaf_capacity(leaf_cap); - wrapper_->set_rtree_variant(variant_str); // load raw WKB data into memory auto field_datas = diff --git a/internal/core/src/index/RTreeIndex.h b/internal/core/src/index/RTreeIndex.h index d47f8ea5182..d8fbec50bf3 100644 --- a/internal/core/src/index/RTreeIndex.h +++ b/internal/core/src/index/RTreeIndex.h @@ -67,7 +67,9 @@ class RTreeIndex : public ScalarIndex { if (is_built_) { return total_num_rows_; } - return wrapper_ ? wrapper_->count() + static_cast(null_offset_.size()) : 0; + return wrapper_ ? wrapper_->count() + + static_cast(null_offset_.size()) + : 0; } // BuildWithRawDataForUT should be only used in ut. Only string is supported. 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 index 3cf499126f0..ec00790f9aa 100644 --- a/internal/core/src/index/RTreeIndexWrapper.cpp +++ b/internal/core/src/index/RTreeIndexWrapper.cpp @@ -10,61 +10,29 @@ // 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 { -// Custom visitor for collecting query results -class GeometryVisitor : public SpatialIndex::IVisitor { - public: - explicit GeometryVisitor(std::vector& results) - : results_(results) { - } - - virtual ~GeometryVisitor() = default; - - void - visitNode(const SpatialIndex::INode& n) override { - // Not needed for our use case - } - - void - visitData(const SpatialIndex::IData& d) override { - // Store the identifier (row offset) in results - results_.push_back(static_cast(d.getIdentifier())); - } - - void - visitData(std::vector& v) override { - for (const auto* data : v) { - results_.push_back(static_cast(data->getIdentifier())); - } - } - - private: - std::vector& results_; -}; - RTreeIndexWrapper::RTreeIndexWrapper(std::string& path, bool is_build_mode) : index_path_(path), is_build_mode_(is_build_mode) { if (is_build_mode_) { - // Create directory if it doesn't exist std::filesystem::path dir_path = std::filesystem::path(path).parent_path(); if (!dir_path.empty()) { std::filesystem::create_directories(dir_path); } - - // Create disk storage manager for building - storage_manager_ = std::shared_ptr( - SpatialIndex::StorageManager::createNewDiskStorageManager(path, - 4096)); + // Start with an empty rtree for dynamic insertions + rtree_ = RTree(); } } @@ -78,20 +46,7 @@ RTreeIndexWrapper::add_geometry(const uint8_t* wkb_data, folly::SharedMutexWritePriority::WriteHolder lock(rtree_mutex_); AssertInfo(is_build_mode_, "Cannot add geometry in load mode"); - // Lazily create the R-Tree for dynamic insertion if not present yet - if (rtree_ == nullptr) { - SpatialIndex::id_type index_id; - rtree_ = std::shared_ptr( - SpatialIndex::RTree::createNewRTree(*storage_manager_, - fill_factor_, - index_capacity_, - leaf_capacity_, - dimension_, - rtree_variant_, - index_id)); - index_id_ = index_id; - LOG_WARN("create rtree index for dynamic insertion"); - } + std::unique_lock guard(rtree_mutex_); // Parse WKB data to OGR geometry OGRGeometry* geom = nullptr; @@ -107,163 +62,62 @@ RTreeIndexWrapper::add_geometry(const uint8_t* wkb_data, double minX, minY, maxX, maxY; get_bounding_box(geom, minX, minY, maxX, maxY); - // Create region for the bounding box - double low[2] = {minX, minY}; - double high[2] = {maxX, maxY}; - SpatialIndex::Region region(low, high, 2); - - // Insert into R-Tree with row_offset as identifier - rtree_->insertData( - 0, nullptr, region, static_cast(row_offset)); + // 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); } -// Internal IDataStream implementation over FieldDataBase (WKB string rows) -namespace { -class BulkLoadDataStream : public SpatialIndex::IDataStream { - public: - BulkLoadDataStream( - const std::vector>& - field_datas, - bool nullable) - : field_datas_(field_datas), nullable_param_(nullable) { - // Compute a cheap upper bound for stream size: sum of row counts - total_rows_ = 0; - for (const auto& fd : field_datas_) { - total_rows_ += static_cast(fd->get_num_rows()); - } - rewind(); - } - - ~BulkLoadDataStream() override = default; - - bool - hasNext() override { - return absolute_offset_ < static_cast(total_rows_); - } - - uint32_t - size() override { - // Return upper bound; actual yielded items may be fewer due to - // null rows or invalid WKB filtered in getNext(). - return static_cast(total_rows_); - } - - void - rewind() override { - batch_index_ = 0; - row_in_batch_ = 0; - absolute_offset_ = 0; - } - - SpatialIndex::IData* - getNext() override { - while (batch_index_ < field_datas_.size()) { - const auto& fd = field_datas_[batch_index_]; - auto n = fd->get_num_rows(); - if (row_in_batch_ >= n) { - ++batch_index_; - row_in_batch_ = 0; - continue; - } +// No IDataStream; bulk-load implemented directly for Boost R-tree - int64_t current_row_in_batch = row_in_batch_; - int64_t current_abs = absolute_offset_; - // advance offsets for next call regardless of validity - ++row_in_batch_; - ++absolute_offset_; +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_); - const bool is_nullable_effective = - nullable_param_ || fd->IsNullable(); - if (is_nullable_effective && !fd->is_valid(current_row_in_batch)) { + 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(current_row_in_batch)); + const auto* wkb_str = + static_cast(fd->RawValue(i)); if (wkb_str == nullptr || wkb_str->empty()) { continue; } - - // Parse WKB using OGR to get envelope OGRGeometry* geom = nullptr; - OGRErr err = OGRGeometryFactory::createFromWkb( + auto err = OGRGeometryFactory::createFromWkb( reinterpret_cast(wkb_str->data()), nullptr, &geom, wkb_str->size()); if (err != OGRERR_NONE || geom == nullptr) { - LOG_WARN( - "BulkLoadDataStream: failed to parse WKB at abs {} (batch " - "{}, row {})", - current_abs, - batch_index_, - current_row_in_batch); continue; } - OGREnvelope env; geom->getEnvelope(&env); OGRGeometryFactory::destroyGeometry(geom); - - double low[2] = {env.MinX, env.MinY}; - double high[2] = {env.MaxX, env.MaxY}; - SpatialIndex::Region region(low, high, 2); - - return new SpatialIndex::RTree::Data( - 0, - nullptr, - region, - static_cast(current_abs)); + Box box(Point(env.MinX, env.MinY), Point(env.MaxX, env.MaxY)); + local_values.emplace_back(box, absolute_offset); } - return nullptr; - } - - private: - const std::vector>& field_datas_; - bool nullable_param_ = false; - size_t total_rows_ = 0; - size_t batch_index_ = 0; - int64_t row_in_batch_ = 0; - int64_t absolute_offset_ = 0; -}; -} // anonymous namespace - -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"); - AssertInfo(storage_manager_ != nullptr, "Storage manager is null"); - AssertInfo(rtree_ == nullptr, - "R-Tree already initialized; bulk load requires a fresh tree"); - - BulkLoadDataStream stream(field_datas, nullable); - SpatialIndex::id_type index_id; - try { - rtree_ = std::shared_ptr( - SpatialIndex::RTree::createAndBulkLoadNewRTree( - SpatialIndex::RTree::BLM_STR, - stream, - *storage_manager_, - fill_factor_, - index_capacity_, - leaf_capacity_, - dimension_, - rtree_variant_, - index_id)); - } catch (const std::exception& e) { - LOG_ERROR("Failed to bulk load R-Tree: {}", e.what()); } - - index_id_ = index_id; - LOG_INFO("R-Tree bulk load completed with {} entries", - rtree_ ? "some" : "none"); + values_.swap(local_values); + rtree_ = RTree(values_.begin(), values_.end()); + LOG_INFO("R-Tree bulk load (Boost) completed with {} entries", + values_.size()); } void @@ -274,6 +128,7 @@ RTreeIndexWrapper::finish() { // 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; @@ -281,51 +136,29 @@ RTreeIndexWrapper::finish() { AssertInfo(is_build_mode_, "Cannot finish in load mode"); - // If rtree_ is already reset, we have nothing left to do. Mark finished - // and return. - if (rtree_ == nullptr) { - LOG_DEBUG( - "RTreeIndexWrapper::finish() called with null rtree_, likely " - "already finished."); - finished_ = true; - return; - } - - // Explicitly flush the index header & buffers to disk to guarantee - // consistency before releasing resources. - rtree_->flush(); - - // NOTE: rtree_ internally holds a pointer to the storage manager. We must - // make sure rtree_ is destroyed BEFORE the storage manager. - - // 1. Release rtree_ first so its destructor can safely write the header - // using a still-valid storage_manager_. - rtree_.reset(); - - // 2. Now it is safe to release the storage manager. - storage_manager_.reset(); - - // 3. Write meta file with index parameters for reliable loading. + // 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; - meta["index_id"] = index_id_; - meta["variant"] = static_cast(rtree_variant_); - meta["fill_factor"] = fill_factor_; - meta["index_capacity"] = index_capacity_; - meta["leaf_capacity"] = leaf_capacity_; + // 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 meta json: {}", e.what()); + LOG_WARN("Failed to write R-Tree files: {}", e.what()); } finished_ = true; - LOG_INFO("R-Tree index finished building and saved to {}", index_path_); + LOG_INFO("R-Tree index (Boost) finished building and saved to {}", + index_path_); } void @@ -335,32 +168,25 @@ RTreeIndexWrapper::load() { AssertInfo(!is_build_mode_, "Cannot load in build mode"); + std::unique_lock guard(rtree_mutex_); try { - // Load storage manager - storage_manager_ = std::shared_ptr( - SpatialIndex::StorageManager::loadDiskStorageManager(index_path_)); - - // Determine index id from meta json if available - SpatialIndex::id_type idx_id_to_load = 0; + // Read meta (optional) try { std::ifstream ifs(index_path_ + ".meta.json"); if (ifs.good()) { auto meta = nlohmann::json::parse(ifs); - if (meta.contains("index_id")) { - idx_id_to_load = - meta["index_id"].get(); - } + // 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, fallback to default id 0: {}", - e.what()); + LOG_WARN("Failed to read meta json: {}", e.what()); } - // Load R-Tree index with the resolved id - rtree_ = std::shared_ptr( - SpatialIndex::RTree::loadRTree(*storage_manager_, idx_id_to_load)); + // Read binary data + RTreeSerializer::loadBinary(rtree_, index_path_ + ".bgi"); - LOG_INFO("R-Tree index loaded from {}", index_path_); + 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 {}: {}", @@ -373,32 +199,25 @@ void RTreeIndexWrapper::query_candidates(proto::plan::GISFunctionFilterExpr_GISOp op, const OGRGeometry& query_geom, std::vector& candidate_offsets) { - // Acquire read lock to protect rtree_ access during query - folly::SharedMutexWritePriority::ReadHolder lock(rtree_mutex_); - - AssertInfo(rtree_ != nullptr, "R-Tree index not initialized"); - 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 region - double low[2] = {minX, minY}; - double high[2] = {maxX, maxY}; - SpatialIndex::Region query_region(low, high, 2); - - // Create visitor for collecting results - GeometryVisitor visitor(candidate_offsets); - - // Perform query based on operation type - switch (op) { - default: - // For all GIS operations, we use intersection query as coarse filtering - // The exact geometric relationship will be checked in the refinement phase - rtree_->intersectsWithQuery(query_region, visitor); - break; + // 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 {}", @@ -425,53 +244,8 @@ RTreeIndexWrapper::get_bounding_box(const OGRGeometry* geom, int64_t RTreeIndexWrapper::count() const { - // Acquire read lock to protect rtree_ access during count operation - folly::SharedMutexWritePriority::ReadHolder lock(rtree_mutex_); - - if (rtree_ == nullptr) { - return 0; - } - - // For R-Tree, we need to count the number of data entries - // This is a simplified implementation - in practice, you might want to - // maintain a separate counter during building - SpatialIndex::IStatistics* stats = nullptr; - rtree_->getStatistics(&stats); - if (stats != nullptr) { - int64_t count = stats->getNumberOfData(); - delete stats; - return count; - } - return 0; + return static_cast(rtree_.size()); } -void -RTreeIndexWrapper::set_rtree_variant(const std::string& variant_str) { - if (variant_str == "RSTAR") { - rtree_variant_ = SpatialIndex::RTree::RV_RSTAR; - } else if (variant_str == "QUADRATIC") { - LOG_WARN("QUADRATIC variant is not supported, using RSTAR instead"); - rtree_variant_ = SpatialIndex::RTree::RV_RSTAR; - } else if (variant_str == "LINEAR") { - rtree_variant_ = SpatialIndex::RTree::RV_LINEAR; - } else { - PanicInfo(ErrorCode::UnexpectedError, - fmt::format("Invalid R-Tree variant: {}", variant_str)); - } -} - -void -RTreeIndexWrapper::set_fill_factor(double fill_factor) { - fill_factor_ = fill_factor; -} - -void -RTreeIndexWrapper::set_index_capacity(uint32_t index_capacity) { - index_capacity_ = index_capacity; -} - -void -RTreeIndexWrapper::set_leaf_capacity(uint32_t leaf_capacity) { - leaf_capacity_ = leaf_capacity; -} +// index/leaf capacity setters removed; not applicable for Boost rtree } // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/RTreeIndexWrapper.h b/internal/core/src/index/RTreeIndexWrapper.h index 9a2f3a201b9..f6240f52b07 100644 --- a/internal/core/src/index/RTreeIndexWrapper.h +++ b/internal/core/src/index/RTreeIndexWrapper.h @@ -12,11 +12,12 @@ #pragma once #include +#include #include #include +#include +#include #include "ogr_geometry.h" -#include "spatialindex/SpatialIndex.h" -#include "common/Types.h" #include "pb/plan.pb.h" #include @@ -27,10 +28,13 @@ class FieldDataBase; namespace milvus::index { +namespace bg = boost::geometry; +namespace bgi = boost::geometry::index; + /** - * @brief Wrapper class for libspatialindex R-Tree functionality + * @brief Wrapper class for boost R-Tree functionality * - * This class provides a simplified interface to libspatialindex library, + * 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. */ @@ -49,10 +53,7 @@ class RTreeIndexWrapper { ~RTreeIndexWrapper(); /** - * @brief Add a geometry to the index - * @param wkb_data Pointer to WKB binary data - * @param len Length of WKB data - * @param row_offset Row offset (used as identifier) + * @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); @@ -99,17 +100,8 @@ class RTreeIndexWrapper { int64_t count() const; - void - set_rtree_variant(const std::string& variant_str); - - void - set_fill_factor(double fill_factor); - - void - set_index_capacity(uint32_t index_capacity); - - void - set_leaf_capacity(uint32_t leaf_capacity); + // Boost rtree does not use index/leaf capacities; keep only fill factor for + // compatibility (no-op currently) private: /** @@ -128,25 +120,25 @@ class RTreeIndexWrapper { double& maxY); private: - std::shared_ptr storage_manager_; - std::shared_ptr rtree_; + // 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; - SpatialIndex::id_type index_id_ = 0; // persisted to meta for reliable load + + // Serialize access to rtree_ + mutable std::shared_mutex rtree_mutex_; // R-Tree parameters - double fill_factor_ = 0.8; - uint32_t index_capacity_ = 50; - uint32_t leaf_capacity_ = 50; uint32_t dimension_ = 2; - SpatialIndex::RTree::RTreeVariant rtree_variant_ = - SpatialIndex::RTree::RV_RSTAR; - - // Thread safety: protects rtree_ and related operations - mutable folly::SharedMutexWritePriority rtree_mutex_; }; } // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/Utils.cpp b/internal/core/src/index/Utils.cpp index 859f8736805..0abef5fbbe6 100644 --- a/internal/core/src/index/Utils.cpp +++ b/internal/core/src/index/Utils.cpp @@ -251,56 +251,6 @@ GetIndexMetaFromConfig(const Config& config) { return index_meta; } -double -GetFillFactorFromConfig(const Config& config) { - auto fill_factor = GetValueFromConfig(config, FILL_FACTOR_KEY); - AssertInfo(fill_factor.has_value(), - "fill factor not exist in index config"); - try { - return (std::stod(fill_factor.value())); - } catch (const std::logic_error& e) { - auto err_message = fmt::format("invalided fill factor:{}, error:{}", - fill_factor.value(), - e.what()); - LOG_ERROR(err_message); - throw std::logic_error(err_message); - } -} - -uint32_t -GetIndexCapacityFromConfig(const Config& config) { - auto index_capacity = - GetValueFromConfig(config, INDEX_CAPACITY_KEY); - AssertInfo(index_capacity.has_value(), - "index capacity not exist in index config"); - try { - return (std::stoi(index_capacity.value())); - } catch (const std::logic_error& e) { - auto err_message = fmt::format("invalided index capacity:{}, error:{}", - index_capacity.value(), - e.what()); - LOG_ERROR(err_message); - throw std::logic_error(err_message); - } -} - -uint32_t -GetLeafCapacityFromConfig(const Config& config) { - auto leaf_capacity = - GetValueFromConfig(config, LEAF_CAPACITY_KEY); - AssertInfo(leaf_capacity.has_value(), - "leaf capacity not exist in index config"); - try { - return (std::stoi(leaf_capacity.value())); - } catch (const std::logic_error& e) { - auto err_message = fmt::format("invalided leaf capacity:{}, error:{}", - leaf_capacity.value(), - e.what()); - LOG_ERROR(err_message); - throw std::logic_error(err_message); - } -} - Config ParseConfigFromIndexParams( const std::map& index_params) { diff --git a/internal/core/src/index/Utils.h b/internal/core/src/index/Utils.h index 13bd84dae6f..97067b48eab 100644 --- a/internal/core/src/index/Utils.h +++ b/internal/core/src/index/Utils.h @@ -163,15 +163,6 @@ GetFieldDataMetaFromConfig(const Config& config); storage::IndexMeta GetIndexMetaFromConfig(const Config& config); -double -GetFillFactorFromConfig(const Config& config); - -uint32_t -GetIndexCapacityFromConfig(const Config& config); - -uint32_t -GetLeafCapacityFromConfig(const Config& config); - Config ParseConfigFromIndexParams( const std::map& index_params); diff --git a/internal/proxy/task_index.go b/internal/proxy/task_index.go index a0ecbcd6c54..63e94e3fd2a 100644 --- a/internal/proxy/task_index.go +++ b/internal/proxy/task_index.go @@ -489,19 +489,6 @@ func checkTrain(ctx context.Context, field *schemapb.FieldSchema, indexParams ma } } - if indexType == indexparamcheck.IndexRTREE { - // Apply default RTree parameters if not provided - rtreeParams := paramtable.Get().AutoIndexConfig.RTreeAutoIndexParams.GetAsJSONMap() - for k, v := range rtreeParams { - if k != common.IndexTypeKey { // Don't override index_type - _, exist := indexParams[k] - if !exist { - indexParams[k] = v - } - } - } - } - 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/constraints.go b/internal/util/indexparamcheck/constraints.go index 64911880221..cf0863d7e19 100644 --- a/internal/util/indexparamcheck/constraints.go +++ b/internal/util/indexparamcheck/constraints.go @@ -49,32 +49,6 @@ const ( BM25B = "bm25_b" MaxBitmapCardinalityLimit = 1000 - - // RTree Index Param - RTreeFillFactor = "fillFactor" - RTreeIndexCapacity = "indexCapacity" - RTreeLeafCapacity = "leafCapacity" - RTreeDim = "dim" - RTreeRV = "rv" - - // RTree parameter constraints - MinRTreeFillFactor = 0.1 - MaxRTreeFillFactor = 1.0 - DefaultRTreeFillFactor = 0.8 - - MinRTreeIndexCapacity = 2 - MaxRTreeIndexCapacity = 1000 - DefaultRTreeIndexCapacity = 128 - - MinRTreeLeafCapacity = 2 - MaxRTreeLeafCapacity = 1000 - DefaultRTreeLeafCapacity = 128 - - MinRTreeDim = 2 - MaxRTreeDim = 2 - DefaultRTreeDim = 2 - - DefaultRTreeRV = "RSTAR" ) var ( diff --git a/internal/util/indexparamcheck/rtree_checker.go b/internal/util/indexparamcheck/rtree_checker.go index 55ea1024618..d7144ad1885 100644 --- a/internal/util/indexparamcheck/rtree_checker.go +++ b/internal/util/indexparamcheck/rtree_checker.go @@ -20,7 +20,6 @@ import ( "fmt" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" - "github.com/milvus-io/milvus/pkg/v2/util/funcutil" "github.com/milvus-io/milvus/pkg/v2/util/typeutil" ) @@ -34,42 +33,6 @@ func (c *RTREEChecker) CheckTrain(dataType schemapb.DataType, params map[string] return fmt.Errorf("RTREE index can only be built on geometry field") } - // Set default values if not provided - setDefaultIfNotExist(params, RTreeFillFactor, fmt.Sprintf("%f", DefaultRTreeFillFactor)) - setDefaultIfNotExist(params, RTreeIndexCapacity, fmt.Sprintf("%d", DefaultRTreeIndexCapacity)) - setDefaultIfNotExist(params, RTreeLeafCapacity, fmt.Sprintf("%d", DefaultRTreeLeafCapacity)) - setDefaultIfNotExist(params, RTreeDim, fmt.Sprintf("%d", DefaultRTreeDim)) - setDefaultIfNotExist(params, RTreeRV, DefaultRTreeRV) - - // Validate fillFactor - if !CheckFloatByRange(params, RTreeFillFactor, MinRTreeFillFactor, MaxRTreeFillFactor) { - return errOutOfRange(params[RTreeFillFactor], MinRTreeFillFactor, MaxRTreeFillFactor) - } - - // Validate indexCapacity - if !CheckIntByRange(params, RTreeIndexCapacity, MinRTreeIndexCapacity, MaxRTreeIndexCapacity) { - return errOutOfRange(params[RTreeIndexCapacity], MinRTreeIndexCapacity, MaxRTreeIndexCapacity) - } - - // Validate leafCapacity - if !CheckIntByRange(params, RTreeLeafCapacity, MinRTreeLeafCapacity, MaxRTreeLeafCapacity) { - return errOutOfRange(params[RTreeLeafCapacity], MinRTreeLeafCapacity, MaxRTreeLeafCapacity) - } - - // Validate dim - if !CheckIntByRange(params, RTreeDim, MinRTreeDim, MaxRTreeDim) { - return errOutOfRange(params[RTreeDim], MinRTreeDim, MaxRTreeDim) - } - - // Validate rv - rvValue, exists := params[RTreeRV] - if exists { - validRVValues := []string{"LINEAR", "QUADRATIC", "RSTAR"} - if !funcutil.SliceContain(validRVValues, rvValue) { - return fmt.Errorf("rv value %s is not supported, supported values: %v", rvValue, validRVValues) - } - } - return c.scalarIndexChecker.CheckTrain(dataType, params) } diff --git a/pkg/util/paramtable/autoindex_param.go b/pkg/util/paramtable/autoindex_param.go index 5bc7faa038d..61c7020cb81 100644 --- a/pkg/util/paramtable/autoindex_param.go +++ b/pkg/util/paramtable/autoindex_param.go @@ -56,7 +56,6 @@ type AutoIndexConfig struct { ScalarJSONIndexType ParamItem `refreshable:"true"` ScalarGeometryIndexType ParamItem `refreshable:"true"` - RTreeAutoIndexParams ParamItem `refreshable:"true"` BitmapCardinalityLimit ParamItem `refreshable:"true"` } @@ -234,13 +233,6 @@ func (p *AutoIndexConfig) init(base *BaseTable) { } p.ScalarGeometryIndexType.Init(base.mgr) - p.RTreeAutoIndexParams = ParamItem{ - Key: "scalarAutoIndex.params.rtree", - Version: "2.5.16", - DefaultValue: `{"fillFactor": 0.8, "indexCapacity": 100, "leafCapacity": 100, "dim": 2, "rv": "RSTAR", "index_type": "RTREE"}`, - } - p.RTreeAutoIndexParams.Init(base.mgr) - p.BitmapCardinalityLimit = ParamItem{ Key: "scalarAutoIndex.params.bitmapCardinalityLimit", Version: "2.5.0",