From 8e1ce15146d7c20341b8223878eed42962915bd4 Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Thu, 21 Aug 2025 14:41:46 +0800 Subject: [PATCH] fix: ngram index is mistakenly used for unsopported operations (#43955) issue: https://github.com/milvus-io/milvus/issues/43917 1. fix ngrma index to be mistakenly used for unsopported operation 2. fix potential uaf problem --------- Signed-off-by: SpadeA --- .../core/src/exec/expression/UnaryExpr.cpp | 420 +++++++++--------- internal/core/unittest/test_ngram_query.cpp | 30 +- 2 files changed, 234 insertions(+), 216 deletions(-) diff --git a/internal/core/src/exec/expression/UnaryExpr.cpp b/internal/core/src/exec/expression/UnaryExpr.cpp index 813fc9ac87..40887c690b 100644 --- a/internal/core/src/exec/expression/UnaryExpr.cpp +++ b/internal/core/src/exec/expression/UnaryExpr.cpp @@ -329,8 +329,9 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplArray(EvalCtx& context) { } int processed_cursor = 0; auto execute_sub_batch = - [ op_type, &processed_cursor, & - bitmap_input ]( + [op_type, + &processed_cursor, + &bitmap_input]( const milvus::ArrayView* data, const bool* valid_data, const int32_t* offsets, @@ -339,185 +340,186 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplArray(EvalCtx& context) { TargetBitmapView valid_res, ValueType val, int index) { - switch (op_type) { - case proto::plan::GreaterThan: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; + switch (op_type) { + case proto::plan::GreaterThan: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::GreaterEqual: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::LessThan: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::LessEqual: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::Equal: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::NotEqual: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::PrefixMatch: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::Match: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::PostfixMatch: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + case proto::plan::InnerMatch: { + UnaryElementFuncForArray + func; + func(data, + valid_data, + size, + val, + index, + res, + valid_res, + bitmap_input, + processed_cursor, + offsets); + break; + } + default: + ThrowInfo( + OpTypeInvalid, + fmt::format( + "unsupported operator type for unary expr: {}", + op_type)); } - case proto::plan::GreaterEqual: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::LessThan: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::LessEqual: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::Equal: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::NotEqual: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::PrefixMatch: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::Match: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::PostfixMatch: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - case proto::plan::InnerMatch: { - UnaryElementFuncForArray - func; - func(data, - valid_data, - size, - val, - index, - res, - valid_res, - bitmap_input, - processed_cursor, - offsets); - break; - } - default: - ThrowInfo( - OpTypeInvalid, - fmt::format("unsupported operator type for unary expr: {}", - op_type)); - } - processed_cursor += size; - }; + processed_cursor += size; + }; int64_t processed_size; if (has_offset_input_) { processed_size = @@ -717,16 +719,18 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplJson(EvalCtx& context) { } while (false) int processed_cursor = 0; - auto execute_sub_batch = - [ op_type, pointer, &processed_cursor, & - bitmap_input ]( - const milvus::Json* data, - const bool* valid_data, - const int32_t* offsets, - const int size, - TargetBitmapView res, - TargetBitmapView valid_res, - ExprValueType val) { + auto execute_sub_batch = [op_type, + pointer, + &processed_cursor, + &bitmap_input]( + const milvus::Json* data, + const bool* valid_data, + const int32_t* offsets, + const int size, + TargetBitmapView res, + TargetBitmapView valid_res, + ExprValueType val) { bool has_bitmap_input = !bitmap_input.empty(); switch (op_type) { case proto::plan::GreaterThan: { @@ -1766,16 +1770,17 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplForData(EvalCtx& context) { auto expr_type = expr_->op_type_; size_t processed_cursor = 0; - auto execute_sub_batch = - [ expr_type, &processed_cursor, & - bitmap_input ]( - const T* data, - const bool* valid_data, - const int32_t* offsets, - const int size, - TargetBitmapView res, - TargetBitmapView valid_res, - IndexInnerType val) { + auto execute_sub_batch = [expr_type, + &processed_cursor, + &bitmap_input]( + const T* data, + const bool* valid_data, + const int32_t* offsets, + const int size, + TargetBitmapView res, + TargetBitmapView valid_res, + IndexInnerType val) { switch (expr_type) { case proto::plan::GreaterThan: { UnaryElementFunc func; @@ -1947,7 +1952,11 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplForData(EvalCtx& context) { template bool PhyUnaryRangeFilterExpr::CanUseIndex() { - use_index_ = is_index_mode_ && SegmentExpr::CanUseIndex(expr_->op_type_); + use_index_ = + is_index_mode_ && SegmentExpr::CanUseIndex(expr_->op_type_) && + // Ngram index should be used in specific execution path (CanExecNgramMatch -> ExecNgramMatch). + // TODO: if multiple indexes are supported, this logic should be changed + !segment_->HasNgramIndex(field_id_); return use_index_; } @@ -2076,15 +2085,14 @@ PhyUnaryRangeFilterExpr::ExecNgramMatch() { } if (cached_ngram_match_res_ == nullptr) { - index::NgramInvertedIndex* index; + PinWrapper pinned_index; if (expr_->column_.data_type_ == DataType::JSON) { - auto pinned_index = segment_->GetNgramIndexForJson( + pinned_index = segment_->GetNgramIndexForJson( field_id_, milvus::Json::pointer(expr_->column_.nested_path_)); - index = pinned_index.get(); } else { - auto pinned_index = segment_->GetNgramIndex(field_id_); - index = pinned_index.get(); + pinned_index = segment_->GetNgramIndex(field_id_); } + index::NgramInvertedIndex* index = pinned_index.get(); AssertInfo(index != nullptr, "ngram index should not be null, field_id: {}", field_id_.get()); diff --git a/internal/core/unittest/test_ngram_query.cpp b/internal/core/unittest/test_ngram_query.cpp index 6d833dcece..d827c55806 100644 --- a/internal/core/unittest/test_ngram_query.cpp +++ b/internal/core/unittest/test_ngram_query.cpp @@ -147,15 +147,16 @@ test_ngram_with_data(const boost::container::vector& data, nb, 8192, 0); - - std::optional bitset_opt = - index->ExecuteQuery(literal, op_type, &segment_expr); - if (forward_to_br) { - ASSERT_TRUE(!bitset_opt.has_value()); - } else { - auto bitset = std::move(bitset_opt.value()); - for (size_t i = 0; i < nb; i++) { - ASSERT_EQ(bitset[i], expected_result[i]); + if (op_type != proto::plan::OpType::Equal) { + std::optional bitset_opt = + index->ExecuteQuery(literal, op_type, &segment_expr); + if (forward_to_br) { + ASSERT_TRUE(!bitset_opt.has_value()); + } else { + auto bitset = std::move(bitset_opt.value()); + for (size_t i = 0; i < nb; i++) { + ASSERT_EQ(bitset[i], expected_result[i]); + } } } } @@ -237,8 +238,13 @@ TEST(NgramIndex, TestNgramWikiEpisode) { // within min-max_gram { + // equal, all should fail + std::vector expected_result{false, false, false, false, false}; + test_ngram_with_data( + data, "ary", proto::plan::OpType::Equal, expected_result); + // inner match - std::vector expected_result{true, true, true, true, true}; + expected_result = {true, true, true, true, true}; test_ngram_with_data( data, "ary", proto::plan::OpType::InnerMatch, expected_result); @@ -412,6 +418,10 @@ TEST(NgramIndex, TestNgramJson) { proto::plan::OpType>> test_cases; proto::plan::GenericValue value; + value.set_string_val("liz"); + test_cases.push_back(std::make_tuple( + value, std::vector{}, proto::plan::OpType::Equal)); + value.set_string_val("nothing"); test_cases.push_back(std::make_tuple( value, std::vector{}, proto::plan::OpType::InnerMatch));