diff --git a/internal/core/src/index/NgramInvertedIndex.cpp b/internal/core/src/index/NgramInvertedIndex.cpp index b967df499c..b8cfc01806 100644 --- a/internal/core/src/index/NgramInvertedIndex.cpp +++ b/internal/core/src/index/NgramInvertedIndex.cpp @@ -18,11 +18,29 @@ namespace milvus::index { +const std::string NGRAM_AVG_ROW_SIZE_FILE_NAME = "ngram_avg_row_size"; + const JsonCastType JSON_CAST_TYPE = JsonCastType::FromString("VARCHAR"); // ngram index doesn't need cast function const JsonCastFunction JSON_CAST_FUNCTION = JsonCastFunction::FromString("unknown"); +constexpr size_t kLargeRowThreshold = 5000; +constexpr size_t kMediumRowThreshold = 1000; +constexpr size_t kSmallRowThreshold = 100; + +constexpr double kPreFilterHitRateThreshold = 0.20; // 20% +// Default avg_row_size for compatibility (older indexes without metadata) +constexpr size_t kDefaultAvgRowSize = kLargeRowThreshold; +// Iterative strategy parameters +constexpr size_t kMaxIterations = 5; +constexpr double kBreakThreshold = 0.002; // 0.2% + +constexpr size_t kMaxIterationsForMediumRow = 3; + +constexpr double kBreakThresholdForSmallRow = 0.01; // 1% +constexpr size_t kMaxIterationsForSmallRow = 2; + // for string/varchar type NgramInvertedIndex::NgramInvertedIndex(const storage::FileManagerContext& ctx, const NgramParams& params) @@ -68,6 +86,25 @@ NgramInvertedIndex::BuildWithFieldData(const std::vector& datas) { if (schema_.data_type() == proto::schema::DataType::JSON) { BuildWithJsonFieldData(datas); } else { + // Calculate avg_row_size for String/VarChar types + size_t total_bytes = 0; + size_t total_rows = 0; + for (const auto& data : datas) { + auto n = data->get_num_rows(); + for (size_t i = 0; i < n; i++) { + if (schema_.nullable() && !data->is_valid(i)) { + continue; + } + auto* str = static_cast(data->RawValue(i)); + if (str) { + total_bytes += str->size(); + total_rows += 1; + } + } + } + avg_row_size_ = total_rows > 0 ? total_bytes / total_rows : 0; + LOG_INFO("Ngram index avg_row_size: {} bytes", avg_row_size_); + InvertedIndexTantivy::BuildWithFieldData(datas); } } @@ -83,6 +120,11 @@ NgramInvertedIndex::BuildWithJsonFieldData( nested_path_); index_build_begin_ = std::chrono::system_clock::now(); + + // Track total bytes and rows for avg_row_size calculation + size_t total_bytes = 0; + size_t total_rows = 0; + ProcessJsonFieldData( field_datas, this->schema_, @@ -90,9 +132,15 @@ NgramInvertedIndex::BuildWithJsonFieldData( JSON_CAST_TYPE, JSON_CAST_FUNCTION, // add data - [this](const std::string* data, int64_t size, int64_t offset) { + [this, &total_bytes, &total_rows]( + const std::string* data, int64_t size, int64_t offset) { this->wrapper_->template add_array_data( data, size, offset); + // Track row size + if (data && size > 0) { + total_bytes += data->size(); + total_rows++; + } }, // handle null [this](int64_t offset) { this->null_offset_.push_back(offset); }, @@ -105,9 +153,25 @@ NgramInvertedIndex::BuildWithJsonFieldData( this->error_recorder_.Record(json, nested_path, error); }); + avg_row_size_ = total_rows > 0 ? total_bytes / total_rows : 0; + LOG_INFO("Ngram index (JSON) avg_row_size: {} bytes", avg_row_size_); + error_recorder_.PrintErrStats(); } +BinarySet +NgramInvertedIndex::Serialize(const Config& config) { + auto res_set = InvertedIndexTantivy::Serialize(config); + + // Serialize avg_row_size + std::shared_ptr avg_row_size_data(new uint8_t[sizeof(size_t)]); + memcpy(avg_row_size_data.get(), &avg_row_size_, sizeof(size_t)); + res_set.Append( + NGRAM_AVG_ROW_SIZE_FILE_NAME, avg_row_size_data, sizeof(size_t)); + + return res_set; +} + IndexStatsPtr NgramInvertedIndex::Upload(const Config& config) { finish(); @@ -117,13 +181,64 @@ NgramInvertedIndex::Upload(const Config& config) { .count(); LOG_INFO( "index build done for ngram index, data type {}, field id: {}, " - "duration: {}s", + "duration: {}s, avg_row_size: {} bytes", schema_.data_type(), field_id_, - index_build_duration); + index_build_duration, + avg_row_size_); + return InvertedIndexTantivy::Upload(config); } +void +NgramInvertedIndex::LoadIndexMetas(const std::vector& index_files, + const Config& config) { + // Call parent to load null_offset + InvertedIndexTantivy::LoadIndexMetas(index_files, config); + + // Load avg_row_size + auto avg_row_size_it = std::find_if( + index_files.begin(), index_files.end(), [](const std::string& file) { + return file.find(NGRAM_AVG_ROW_SIZE_FILE_NAME) != std::string::npos; + }); + + if (avg_row_size_it != index_files.end()) { + auto load_priority = + GetValueFromConfig( + config, milvus::LOAD_PRIORITY) + .value_or(milvus::proto::common::LoadPriority::HIGH); + // avg_row_size is only 8 bytes, never sliced + auto index_datas = mem_file_manager_->LoadIndexToMemory( + {*avg_row_size_it}, load_priority); + auto avg_row_size_data = + std::move(index_datas.at(NGRAM_AVG_ROW_SIZE_FILE_NAME)); + memcpy( + &avg_row_size_, avg_row_size_data->PayloadData(), sizeof(size_t)); + LOG_INFO("Loaded ngram index avg_row_size: {} bytes", avg_row_size_); + } else { + avg_row_size_ = kDefaultAvgRowSize; + LOG_INFO("No avg_row_size metadata found, using default: {}", + kDefaultAvgRowSize); + } +} + +void +NgramInvertedIndex::RetainTantivyIndexFiles( + std::vector& index_files) { + // Call parent to filter null_offset + InvertedIndexTantivy::RetainTantivyIndexFiles(index_files); + + // Also filter avg_row_size + index_files.erase( + std::remove_if(index_files.begin(), + index_files.end(), + [](const std::string& file) { + return file.find(NGRAM_AVG_ROW_SIZE_FILE_NAME) != + std::string::npos; + }), + index_files.end()); +} + void NgramInvertedIndex::Load(milvus::tracer::TraceContext ctx, const Config& config) { @@ -131,38 +246,19 @@ NgramInvertedIndex::Load(milvus::tracer::TraceContext ctx, GetValueFromConfig>(config, INDEX_FILES); AssertInfo(index_files.has_value(), "index file paths is empty when load ngram index"); + auto files_value = index_files.value(); + + LoadIndexMetas(files_value, config); + RetainTantivyIndexFiles(files_value); auto load_priority = GetValueFromConfig( config, milvus::LOAD_PRIORITY) .value_or(milvus::proto::common::LoadPriority::HIGH); - auto files_value = index_files.value(); - auto it = std::find_if( - files_value.begin(), files_value.end(), [](const std::string& file) { - constexpr std::string_view suffix{"/index_null_offset"}; - return file.size() >= suffix.size() && - std::equal(suffix.rbegin(), suffix.rend(), file.rbegin()); - }); - if (it != files_value.end()) { - std::vector file; - file.push_back(*it); - files_value.erase(it); - auto index_datas = - mem_file_manager_->LoadIndexToMemory(file, load_priority); - BinarySet binary_set; - AssembleIndexDatas(index_datas, binary_set); - // clear index_datas to free memory early - index_datas.clear(); - auto index_valid_data = binary_set.GetByName("index_null_offset"); - null_offset_.resize((size_t)index_valid_data->size / sizeof(size_t)); - memcpy(null_offset_.data(), - index_valid_data->data.get(), - (size_t)index_valid_data->size); - } - disk_file_manager_->CacheNgramIndexToDisk(files_value, load_priority); AssertInfo( tantivy_index_exist(path_.c_str()), "index not exist: {}", path_); + auto load_in_mmap = GetValueFromConfig(config, ENABLE_MMAP).value_or(true); wrapper_ = std::make_shared( @@ -190,7 +286,6 @@ NgramInvertedIndex::ExecuteQuery(const std::string& literal, if (literal.length() < min_gram_) { return std::nullopt; } - if (Count() == 0) { return TargetBitmap{}; } @@ -319,6 +414,48 @@ handle_batch(const T* data, } } +void +NgramInvertedIndex::ApplyIterativeNgramFilter( + const std::vector& sorted_terms, + size_t total_count, + TargetBitmap& bitset) { + auto max_iterations = kMaxIterations; + if (avg_row_size_ < kSmallRowThreshold) { + max_iterations = kMaxIterationsForSmallRow; + } else if (avg_row_size_ < kMediumRowThreshold) { + max_iterations = kMaxIterationsForMediumRow; + } + + for (size_t i = 0; i < std::min(sorted_terms.size(), max_iterations); i++) { + TargetBitmap term_bitset{total_count}; + wrapper_->ngram_term_posting_list(sorted_terms[i], &term_bitset); + bitset &= term_bitset; + + double current_hit_rate = 1.0 * bitset.count() / total_count; + if (current_hit_rate < kBreakThreshold) { + break; + } + if (avg_row_size_ < kSmallRowThreshold && + current_hit_rate < kBreakThresholdForSmallRow) { + break; + } + } +} + +// Strategy selection based on avg_row_size and pre_filter_hit_rate: +// - Batch strategy: query all ngram terms at once via ngram_match_query. +// Used for large rows (>= 5KB) where full matching is efficient, or for +// medium rows (>= 1KB) with high pre_filter_hit_rate (> 20%). +// - Iterative strategy: query terms one by one, sorted by doc_freq (rarest first). +// Used for small/medium rows with low pre_filter_hit_rate, allowing early +// termination when hit_rate drops below threshold. +bool +NgramInvertedIndex::ShouldUseBatchStrategy(double pre_filter_hit_rate) const { + return avg_row_size_ >= kLargeRowThreshold || + (avg_row_size_ >= kMediumRowThreshold && + pre_filter_hit_rate > kPreFilterHitRateThreshold); +} + template std::optional NgramInvertedIndex::ExecuteQueryWithPredicate(const std::string& literal, @@ -327,24 +464,43 @@ NgramInvertedIndex::ExecuteQueryWithPredicate(const std::string& literal, bool need_post_filter, const TargetBitmap* pre_filter) { auto total_count = static_cast(Count()); - // Phase 1: ngram index filtering - TargetBitmap bitset{total_count}; - wrapper_->ngram_match_query(literal, min_gram_, max_gram_, &bitset); - auto phase1_hit_count = bitset.count(); - if (phase1_hit_count == 0) { + + AssertInfo(total_count > 0, "total_count should be greater than 0"); + + // Calculate pre_filter stats for strategy selection + size_t candidate_count = (pre_filter != nullptr && !pre_filter->empty()) + ? pre_filter->count() + : total_count; + double pre_filter_hit_rate = 1.0 * candidate_count / total_count; + + // Phase 1: ngram index query with adaptive strategy + TargetBitmap bitset(total_count, true); + if (pre_filter != nullptr && !pre_filter->empty()) { + bitset &= *pre_filter; + } + + if (ShouldUseBatchStrategy(pre_filter_hit_rate)) { + TargetBitmap ngram_bitset{total_count}; + wrapper_->ngram_match_query( + literal, min_gram_, max_gram_, &ngram_bitset); + bitset &= ngram_bitset; + } else { + std::vector literals_vec = {literal}; + auto sorted_terms = + wrapper_->ngram_tokenize(literals_vec, min_gram_, max_gram_); + AssertInfo(!sorted_terms.empty(), + "ngram_tokenize should not return empty for valid literal"); + ApplyIterativeNgramFilter(sorted_terms, total_count, bitset); + } + + if (bitset.none()) { return std::move(bitset); } - // Apply pre_filter from previous expressions to reduce phase 2 workload - if (pre_filter != nullptr && !pre_filter->empty()) { - bitset &= *pre_filter; - if (bitset.none()) { - return std::move(bitset); - } - } - auto phase2_input_count = bitset.count(); + auto after_pre_filter_count = bitset.count(); + auto final_result_count = after_pre_filter_count; - // Phase 2: post-filter with predicate + // Phase 2: post-filter if (need_post_filter) { TargetBitmapView res(bitset); @@ -355,22 +511,20 @@ NgramInvertedIndex::ExecuteQueryWithPredicate(const std::string& literal, }; segment->template ProcessAllDataChunkBatched(execute_batch, res); - } + final_result_count = bitset.count(); + } if (auto root_span = tracer::GetRootSpan()) { + double hit_rate_before_phase2 = + 1.0 * after_pre_filter_count / total_count; + double final_hit_rate = 1.0 * final_result_count / total_count; + root_span->SetAttribute("need_post_filter", need_post_filter); - root_span->SetAttribute( - "phase1_hit_rate", - static_cast(phase1_hit_count) / total_count); - root_span->SetAttribute( - "pre_filter_rate", - static_cast(phase2_input_count) / phase1_hit_count); - root_span->SetAttribute( - "phase2_input_rate", - static_cast(phase2_input_count) / total_count); - root_span->SetAttribute( - "final_hit_rate", - static_cast(bitset.count()) / total_count); + root_span->SetAttribute("pre_filter_hit_rate", pre_filter_hit_rate); + root_span->SetAttribute("hit_rate_before_phase2", + hit_rate_before_phase2); + root_span->SetAttribute("final_hit_rate", final_hit_rate); + root_span->SetAttribute("total_count", static_cast(total_count)); } return std::optional(std::move(bitset)); @@ -418,86 +572,99 @@ NgramInvertedIndex::MatchQuery(const std::string& literal, } auto total_count = static_cast(Count()); + AssertInfo(total_count > 0, "total_count should be greater than 0"); - // Phase 1: ngram index filtering - TargetBitmap bitset(total_count, true); auto literals = split_by_wildcard(literal); for (const auto& l : literals) { if (l.length() < min_gram_) { return std::nullopt; } - TargetBitmap tmp_bitset(total_count, false); - wrapper_->ngram_match_query(l, min_gram_, max_gram_, &tmp_bitset); - bitset &= tmp_bitset; } - auto phase1_hit_count = bitset.count(); - if (phase1_hit_count == 0) { + + // Calculate pre_filter stats for strategy selection + size_t candidate_count = (pre_filter != nullptr && !pre_filter->empty()) + ? pre_filter->count() + : total_count; + double pre_filter_hit_rate = 1.0 * candidate_count / total_count; + + TargetBitmap bitset(total_count, true); + if (pre_filter != nullptr && !pre_filter->empty()) { + bitset &= *pre_filter; + } + + if (ShouldUseBatchStrategy(pre_filter_hit_rate)) { + for (const auto& l : literals) { + TargetBitmap tmp_bitset{total_count}; + wrapper_->ngram_match_query(l, min_gram_, max_gram_, &tmp_bitset); + bitset &= tmp_bitset; + } + } else { + auto sorted_terms = + wrapper_->ngram_tokenize(literals, min_gram_, max_gram_); + AssertInfo(!sorted_terms.empty(), + "ngram_tokenize returned empty sorted_terms for iterative " + "strategy"); + ApplyIterativeNgramFilter(sorted_terms, total_count, bitset); + } + + if (bitset.none()) { return std::move(bitset); } - // Apply pre_filter from previous expressions to reduce phase 2 workload - if (pre_filter != nullptr && !pre_filter->empty()) { - bitset &= *pre_filter; - if (bitset.none()) { - return std::move(bitset); - } + auto after_pre_filter_count = bitset.count(); + + TargetBitmapView res(bitset); + + PatternMatchTranslator translator; + auto regex_pattern = translator(literal); + RegexMatcher matcher(regex_pattern); + + if (schema_.data_type() == proto::schema::DataType::JSON) { + auto predicate = [&matcher, this](const milvus::Json& data) { + auto x = data.template at(this->nested_path_); + if (x.error()) { + return false; + } + return matcher(x.value()); + }; + + auto execute_batch = [&predicate](const milvus::Json* data, + const int64_t size, + TargetBitmapView res) { + handle_batch(data, size, res, predicate); + }; + + segment->template ProcessAllDataChunkBatched( + execute_batch, res); + } else { + auto predicate = [&matcher](const std::string_view& data) { + return matcher(data); + }; + + auto execute_batch = [&predicate](const std::string_view* data, + const int64_t size, + TargetBitmapView res) { + handle_batch(data, size, res, predicate); + }; + + segment->template ProcessAllDataChunkBatched( + execute_batch, res); } - auto phase2_input_count = bitset.count(); - // Phase 2: post-filter with regex - { - TargetBitmapView res(bitset); - - PatternMatchTranslator translator; - auto regex_pattern = translator(literal); - RegexMatcher matcher(regex_pattern); - - if (schema_.data_type() == proto::schema::DataType::JSON) { - auto predicate = [&matcher, this](const milvus::Json& data) { - auto x = data.template at(this->nested_path_); - if (x.error()) { - return false; - } - return matcher(x.value()); - }; - - auto execute_batch = [&predicate](const milvus::Json* data, - const int64_t size, - TargetBitmapView res) { - handle_batch(data, size, res, predicate); - }; - - segment->template ProcessAllDataChunkBatched( - execute_batch, res); - } else { - auto predicate = [&matcher](const std::string_view& data) { - return matcher(data); - }; - - auto execute_batch = [&predicate](const std::string_view* data, - const int64_t size, - TargetBitmapView res) { - handle_batch(data, size, res, predicate); - }; - - segment->template ProcessAllDataChunkBatched( - execute_batch, res); - } - } + auto final_result_count = bitset.count(); if (auto root_span = tracer::GetRootSpan()) { - root_span->SetAttribute( - "match_phase1_hit_rate", - static_cast(phase1_hit_count) / total_count); - root_span->SetAttribute( - "match_pre_filter_rate", - static_cast(phase2_input_count) / phase1_hit_count); - root_span->SetAttribute( - "match_phase2_input_rate", - static_cast(phase2_input_count) / total_count); - root_span->SetAttribute( - "match_final_hit_rate", - static_cast(bitset.count()) / total_count); + double hit_rate_before_phase2 = + 1.0 * after_pre_filter_count / total_count; + double final_hit_rate = 1.0 * final_result_count / total_count; + + root_span->SetAttribute("match_pre_filter_hit_rate", + pre_filter_hit_rate); + root_span->SetAttribute("match_hit_rate_before_phase2", + hit_rate_before_phase2); + root_span->SetAttribute("match_final_hit_rate", final_hit_rate); + root_span->SetAttribute("match_total_count", + static_cast(total_count)); } return std::optional(std::move(bitset)); diff --git a/internal/core/src/index/NgramInvertedIndex.h b/internal/core/src/index/NgramInvertedIndex.h index f9656af8e1..d256e0e855 100644 --- a/internal/core/src/index/NgramInvertedIndex.h +++ b/internal/core/src/index/NgramInvertedIndex.h @@ -21,6 +21,7 @@ class SegmentExpr; } // namespace milvus::exec namespace milvus::index { + class NgramInvertedIndex : public InvertedIndexTantivy { public: // for string/varchar type @@ -32,12 +33,22 @@ class NgramInvertedIndex : public InvertedIndexTantivy { const NgramParams& params, const std::string& nested_path); + BinarySet + Serialize(const Config& config) override; + IndexStatsPtr Upload(const Config& config = {}) override; void Load(milvus::tracer::TraceContext ctx, const Config& config) override; + void + LoadIndexMetas(const std::vector& index_files, + const Config& config) override; + + void + RetainTantivyIndexFiles(std::vector& index_files) override; + void BuildWithFieldData(const std::vector& datas) override; @@ -80,10 +91,19 @@ class NgramInvertedIndex : public InvertedIndexTantivy { exec::SegmentExpr* segment, const TargetBitmap* pre_filter); + void + ApplyIterativeNgramFilter(const std::vector& sorted_terms, + size_t total_count, + TargetBitmap& bitset); + + bool + ShouldUseBatchStrategy(double pre_filter_hit_rate) const; + private: uintptr_t min_gram_{0}; uintptr_t max_gram_{0}; int64_t field_id_{0}; + size_t avg_row_size_{0}; std::chrono::time_point index_build_begin_; // for json type diff --git a/internal/core/thirdparty/tantivy/tantivy-binding/include/tantivy-binding.h b/internal/core/thirdparty/tantivy/tantivy-binding/include/tantivy-binding.h index 89dbd1c50c..2d21106cd7 100644 --- a/internal/core/thirdparty/tantivy/tantivy-binding/include/tantivy-binding.h +++ b/internal/core/thirdparty/tantivy/tantivy-binding/include/tantivy-binding.h @@ -27,11 +27,19 @@ struct RustArrayI64 { size_t cap; }; +/// Array of C strings (char*) for returning Vec to C++ +struct RustStringArray { + char **array; + size_t len; + size_t cap; +}; + struct Value { enum class Tag { None, RustArray, RustArrayI64, + RustStringArray, U32, U64, Ptr, @@ -49,6 +57,10 @@ struct Value { RustArrayI64 _0; }; + struct RustStringArray_Body { + RustStringArray _0; + }; + struct U32_Body { uint32_t _0; }; @@ -66,6 +78,7 @@ struct Value { None_Body none; RustArray_Body rust_array; RustArrayI64_Body rust_array_i64; + RustStringArray_Body rust_string_array; U32_Body u32; U64_Body u64; Ptr_Body ptr; @@ -94,6 +107,8 @@ void free_rust_array(RustArray array); void free_rust_array_i64(RustArrayI64 array); +void free_rust_string_array(RustStringArray array); + void free_rust_result(RustResult result); void free_rust_error(const char *error); @@ -293,6 +308,14 @@ RustResult tantivy_ngram_match_query(void *ptr, uintptr_t max_gram, void *bitset); +RustResult tantivy_ngram_tokenize(void *ptr, + const char *const *literals, + uintptr_t literals_len, + uintptr_t min_gram, + uintptr_t max_gram); + +RustResult tantivy_ngram_term_posting_list(void *ptr, const char *term, void *bitset); + RustResult tantivy_match_query(void *ptr, const char *query, uintptr_t min_should_match, diff --git a/internal/core/thirdparty/tantivy/tantivy-binding/src/array.rs b/internal/core/thirdparty/tantivy/tantivy-binding/src/array.rs index 956e3ff46a..805444bf7e 100644 --- a/internal/core/thirdparty/tantivy/tantivy-binding/src/array.rs +++ b/internal/core/thirdparty/tantivy/tantivy-binding/src/array.rs @@ -98,12 +98,76 @@ pub extern "C" fn free_rust_array_i64(array: RustArrayI64) { } } +/// Array of C strings (char*) for returning Vec to C++ +#[repr(C)] +pub struct RustStringArray { + pub array: *mut *mut c_char, + pub len: size_t, + pub cap: size_t, +} + +impl RustStringArray { + pub fn from_vec(vec: Vec) -> RustStringArray { + let len = vec.len(); + let cap = vec.capacity(); + + // Convert Vec to Vec<*mut c_char> + let c_strings: Vec<*mut c_char> = vec + .into_iter() + .map(|s| create_string(&s) as *mut c_char) + .collect(); + + let c_len = c_strings.len(); + let c_cap = c_strings.capacity(); + let ptr = c_strings.leak().as_mut_ptr(); + + RustStringArray { + array: ptr, + len: c_len, + cap: c_cap, + } + } +} + +impl std::default::Default for RustStringArray { + fn default() -> Self { + RustStringArray { + array: std::ptr::null_mut(), + len: 0, + cap: 0, + } + } +} + +impl From> for RustStringArray { + fn from(vec: Vec) -> Self { + RustStringArray::from_vec(vec) + } +} + +#[no_mangle] +pub extern "C" fn free_rust_string_array(array: RustStringArray) { + let RustStringArray { array, len, cap } = array; + if array.is_null() { + return; + } + unsafe { + let vec = Vec::from_raw_parts(array, len, cap); + for s in vec { + if !s.is_null() { + free_rust_string(s); + } + } + } +} + #[allow(dead_code)] #[repr(C)] pub enum Value { None(()), RustArray(RustArray), RustArrayI64(RustArrayI64), + RustStringArray(RustStringArray), U32(u32), U64(u64), Ptr(*mut c_void), @@ -121,7 +185,7 @@ macro_rules! impl_from_for_enum { }; } -impl_from_for_enum!(Value, None => (), RustArrayI64 => RustArrayI64, RustArrayI64 => Vec, RustArray => RustArray, RustArray => Vec, U32 => u32, U64 => u64, Ptr => *mut c_void); +impl_from_for_enum!(Value, None => (), RustArrayI64 => RustArrayI64, RustArrayI64 => Vec, RustArray => RustArray, RustArray => Vec, RustStringArray => RustStringArray, RustStringArray => Vec, U32 => u32, U64 => u64, Ptr => *mut c_void); #[repr(C)] pub struct RustResult { @@ -197,6 +261,11 @@ pub extern "C" fn free_rust_result(result: RustResult) { free_rust_array_i64(array); } } + Value::RustStringArray(array) => { + if !array.array.is_null() { + free_rust_string_array(array); + } + } _ => {} } if !result.error.is_null() { diff --git a/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader.rs b/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader.rs index 25c1bea062..480f686051 100644 --- a/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader.rs +++ b/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader.rs @@ -590,6 +590,87 @@ impl IndexReaderWrapper { let query = BooleanQuery::intersection(term_queries); self.search(&query, bitset) } + + /// Tokenize literals into ngram terms and return them sorted by doc_freq (ascending). + /// + /// For Match type queries like `%xxx%yyy%`, literals = ["xxx", "yyy"]. + /// For InnerMatch/PrefixMatch/PostfixMatch type queries like `%xxx%`, literals = ["xxx"]. + /// + /// Returns: Vec of term strings sorted by doc_freq ascending (rarest first). + pub fn ngram_tokenize( + &self, + literals: &[&str], + min_gram: usize, + max_gram: usize, + ) -> Result> { + // Collect (term_text, Term) pairs to track text alongside terms + let mut all_term_pairs: Vec<(String, Term)> = vec![]; + let mut tokenizer = NgramTokenizer::new(max_gram, max_gram, false).unwrap(); + + for literal in literals { + assert!( + literal.chars().count() >= min_gram, + "literal '{}' must be >= min_gram {}", + literal, + min_gram + ); + + if literal.chars().count() <= max_gram { + all_term_pairs.push(( + literal.to_string(), + Term::from_field_text(self.field, literal), + )); + } else { + let mut token_stream = tokenizer.token_stream(literal); + token_stream.process(&mut |token| { + all_term_pairs.push(( + token.text.clone(), + Term::from_field_text(self.field, &token.text), + )); + }); + } + } + + assert!( + !all_term_pairs.is_empty(), + "ngram_tokenize should not produce empty terms for valid literals" + ); + + // Get doc_freq for each term and sort + let searcher = self.reader.searcher(); + let mut term_with_freq: Vec<(String, u64)> = Vec::with_capacity(all_term_pairs.len()); + + for (text, term) in all_term_pairs.iter() { + let mut total_doc_freq: u64 = 0; + for segment_reader in searcher.segment_readers() { + let inv_index = segment_reader.inverted_index(term.field())?; + total_doc_freq += inv_index.doc_freq(term)? as u64; + } + term_with_freq.push((text.clone(), total_doc_freq)); + } + + // Sort by doc_freq ascending (rarest first) + term_with_freq.sort_by_key(|(_, freq)| *freq); + + // Remove duplicates + let mut seen = std::collections::HashSet::new(); + Ok(term_with_freq + .into_iter() + .filter_map(|(text, _)| { + if seen.insert(text.clone()) { + Some(text) + } else { + None + } + }) + .collect()) + } + + /// Get the posting list for a single ngram term. + /// Sets bits in the bitset for all documents containing the term. + pub fn ngram_term_posting_list(&self, term_str: &str, bitset: *mut c_void) -> Result<()> { + self.term_query_keyword(term_str, bitset) + } } #[cfg(test)] diff --git a/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader_c.rs b/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader_c.rs index f3efff674d..1f7d6efe56 100644 --- a/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader_c.rs +++ b/internal/core/thirdparty/tantivy/tantivy-binding/src/index_reader_c.rs @@ -1,4 +1,3 @@ -use core::slice; use std::ffi::{c_char, c_void, CStr}; use crate::{ @@ -551,3 +550,37 @@ pub extern "C" fn tantivy_ngram_match_query( .into() } } + +#[no_mangle] +pub extern "C" fn tantivy_ngram_tokenize( + ptr: *mut c_void, + literals: *const *const c_char, + literals_len: usize, + min_gram: usize, + max_gram: usize, +) -> RustResult { + let real = ptr as *mut IndexReaderWrapper; + let literals_slice = unsafe { convert_to_rust_slice!(literals, literals_len) }; + + let mut literal_strs: Vec<&str> = Vec::with_capacity(literals_len); + for &lit in literals_slice { + literal_strs.push(cstr_to_str!(lit)); + } + + unsafe { + (*real) + .ngram_tokenize(&literal_strs, min_gram, max_gram) + .into() + } +} + +#[no_mangle] +pub extern "C" fn tantivy_ngram_term_posting_list( + ptr: *mut c_void, + term: *const c_char, + bitset: *mut c_void, +) -> RustResult { + let real = ptr as *mut IndexReaderWrapper; + let term = cstr_to_str!(term); + unsafe { (*real).ngram_term_posting_list(term, bitset).into() } +} diff --git a/internal/core/thirdparty/tantivy/tantivy-wrapper.h b/internal/core/thirdparty/tantivy/tantivy-wrapper.h index 3fe0a8f565..d154e2a660 100644 --- a/internal/core/thirdparty/tantivy/tantivy-wrapper.h +++ b/internal/core/thirdparty/tantivy/tantivy-wrapper.h @@ -1008,6 +1008,52 @@ struct TantivyIndexWrapper { "TantivyIndexWrapper.ngram_match_query: invalid result type"); } + // Tokenize literals into ngram terms and return them sorted by doc_freq (ascending). + // For Match type queries like `%xxx%yyy%`, literals = ["xxx", "yyy"]. + // For InnerMatch type queries like `%xxx%`, literals = ["xxx"]. + std::vector + ngram_tokenize(const std::vector& literals, + uintptr_t min_gram, + uintptr_t max_gram) { + std::vector c_literals; + c_literals.reserve(literals.size()); + for (const auto& lit : literals) { + c_literals.push_back(lit.c_str()); + } + + auto array = tantivy_ngram_tokenize( + reader_, c_literals.data(), c_literals.size(), min_gram, max_gram); + auto res = RustResultWrapper(array); + AssertInfo(res.result_->success, + "TantivyIndexWrapper.ngram_tokenize: {}", + res.result_->error); + AssertInfo(res.result_->value.tag == Value::Tag::RustStringArray, + "TantivyIndexWrapper.ngram_tokenize: invalid result type"); + + // Convert RustStringArray to std::vector + auto& rust_array = res.result_->value.rust_string_array._0; + std::vector result; + result.reserve(rust_array.len); + for (size_t i = 0; i < rust_array.len; i++) { + result.emplace_back(rust_array.array[i]); + } + return result; + } + + // Get the posting list for a single ngram term. + void + ngram_term_posting_list(const std::string& term, void* bitset) { + auto array = + tantivy_ngram_term_posting_list(reader_, term.c_str(), bitset); + auto res = RustResultWrapper(array); + AssertInfo(res.result_->success, + "TantivyIndexWrapper.ngram_term_posting_list: {}", + res.result_->error); + AssertInfo( + res.result_->value.tag == Value::Tag::None, + "TantivyIndexWrapper.ngram_term_posting_list: invalid result type"); + } + // json query template void