mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
enhance: batch processing for ngram (#46648)
issue: https://github.com/milvus-io/milvus/issues/42053 Process ngram in batch rather than all by once. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Batch Processing for N-gram Queries **Core Invariant:** All data iteration is now driven by `batch_size_` as the fundamental unit; for sealed chunked segments processing string/JSON data, processing is strictly stateless to allow specialized batched algorithms. **Simplified Logic:** - Removed the `process_all_chunks` boolean flag from `ProcessMultipleChunksCommon` (renamed to `ProcessDataChunksForMultipleChunk`) as it was redundant—all iteration paths now converge on the same `batch_size_`-driven chunking strategy with unified data size clamping (`std::min(chunk_size, batch_size_ - processed_size)`). - Eliminated wrapper delegation methods (`ProcessDataChunksForMultipleChunk` and `ProcessAllChunksForMultipleChunk` old wrappers) that pointed to a single common implementation with a conditional flag. **No Data Loss or Behavior Regression:** - The new `ProcessAllDataChunkBatched<T>` is an additional stateless public path (requires sealed + chunked segments, type constraints: `std::string_view|Json|ArrayView`) that iterates all `num_data_chunk_` chunks in `batch_size_` granularity without mutating cursor state (`current_data_chunk_`, `current_data_chunk_pos_`), ensuring deterministic re-entrant processing. - Existing cursor-based APIs (`ProcessDataChunksForMultipleChunk`, `ProcessChunkForSealedSeg`) remain unchanged for standard expression evaluation—no segment state is corrupted. - N-gram query execution now routes through `ExecuteQueryWithPredicate<T, Predicate>(literal, segment, predicate, need_post_filter)` which forwards generic predicates and delegates to `segment->ProcessAllDataChunkBatched<T>(execute_batch, res)` for post-filtering, avoiding per-chunk single-pass traversal. **Enhancement:** Generic predicate template `template <typename T, typename Predicate>` with perfect forwarding (`Predicate&& predicate`) replaces the fixed `std::function<bool(const T&)>` signature, eliminating function wrapper overhead for n-gram matcher closures and enabling efficient batch processing callbacks. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
@@ -968,46 +968,37 @@ class SegmentExpr : public Expr {
|
||||
return processed_size;
|
||||
}
|
||||
|
||||
// If process_all_chunks is true, all chunks will be processed and no inner state will be changed.
|
||||
template <typename T,
|
||||
bool NeedSegmentOffsets = false,
|
||||
typename FUNC,
|
||||
typename... ValTypes>
|
||||
int64_t
|
||||
ProcessMultipleChunksCommon(
|
||||
ProcessDataChunksForMultipleChunk(
|
||||
FUNC func,
|
||||
std::function<bool(const milvus::SkipIndex&, FieldId, int)> skip_func,
|
||||
TargetBitmapView res,
|
||||
TargetBitmapView valid_res,
|
||||
bool process_all_chunks,
|
||||
const ValTypes&... values) {
|
||||
int64_t processed_size = 0;
|
||||
|
||||
size_t start_chunk = process_all_chunks ? 0 : current_data_chunk_;
|
||||
|
||||
// prefetch chunks to reduce cache miss latency
|
||||
if (!prefetched_) {
|
||||
std::vector<int64_t> pf_chunk_ids;
|
||||
pf_chunk_ids.reserve(num_data_chunk_ - start_chunk);
|
||||
for (size_t i = start_chunk; i < num_data_chunk_; i++) {
|
||||
pf_chunk_ids.reserve(num_data_chunk_ - current_data_chunk_);
|
||||
for (size_t i = current_data_chunk_; i < num_data_chunk_; i++) {
|
||||
pf_chunk_ids.push_back(i);
|
||||
}
|
||||
segment_->prefetch_chunks(op_ctx_, field_id_, pf_chunk_ids);
|
||||
prefetched_ = true;
|
||||
}
|
||||
|
||||
for (size_t i = start_chunk; i < num_data_chunk_; i++) {
|
||||
for (size_t i = current_data_chunk_; i < num_data_chunk_; i++) {
|
||||
auto data_pos =
|
||||
process_all_chunks
|
||||
? 0
|
||||
: (i == current_data_chunk_ ? current_data_chunk_pos_ : 0);
|
||||
i == current_data_chunk_ ? current_data_chunk_pos_ : 0;
|
||||
|
||||
// if segment is chunked, type won't be growing
|
||||
int64_t size = segment_->chunk_size(field_id_, i) - data_pos;
|
||||
// process a whole chunk if process_all_chunks is true
|
||||
if (!process_all_chunks) {
|
||||
size = std::min(size, batch_size_ - processed_size);
|
||||
}
|
||||
size = std::min(size, batch_size_ - processed_size);
|
||||
|
||||
if (size == 0)
|
||||
continue; //do not go empty-loop at the bound of the chunk
|
||||
@@ -1136,7 +1127,7 @@ class SegmentExpr : public Expr {
|
||||
|
||||
processed_size += size;
|
||||
|
||||
if (!process_all_chunks && processed_size >= batch_size_) {
|
||||
if (processed_size >= batch_size_) {
|
||||
current_data_chunk_ = i;
|
||||
current_data_chunk_pos_ = data_pos + size;
|
||||
break;
|
||||
@@ -1146,33 +1137,6 @@ class SegmentExpr : public Expr {
|
||||
return processed_size;
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
bool NeedSegmentOffsets = false,
|
||||
typename FUNC,
|
||||
typename... ValTypes>
|
||||
int64_t
|
||||
ProcessDataChunksForMultipleChunk(
|
||||
FUNC func,
|
||||
std::function<bool(const milvus::SkipIndex&, FieldId, int)> skip_func,
|
||||
TargetBitmapView res,
|
||||
TargetBitmapView valid_res,
|
||||
const ValTypes&... values) {
|
||||
return ProcessMultipleChunksCommon<T, NeedSegmentOffsets>(
|
||||
func, skip_func, res, valid_res, false, values...);
|
||||
}
|
||||
|
||||
template <typename T, typename FUNC, typename... ValTypes>
|
||||
int64_t
|
||||
ProcessAllChunksForMultipleChunk(
|
||||
FUNC func,
|
||||
std::function<bool(const milvus::SkipIndex&, FieldId, int)> skip_func,
|
||||
TargetBitmapView res,
|
||||
TargetBitmapView valid_res,
|
||||
const ValTypes&... values) {
|
||||
return ProcessMultipleChunksCommon<T>(
|
||||
func, skip_func, res, valid_res, true, values...);
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
bool NeedSegmentOffsets = false,
|
||||
typename FUNC,
|
||||
@@ -1193,20 +1157,45 @@ class SegmentExpr : public Expr {
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename FUNC, typename... ValTypes>
|
||||
// Specialized method for ngram post-filter: processes all data in batches
|
||||
// - Starts from position 0
|
||||
// - Does NOT modify segment state variables (current_data_chunk_, etc.)
|
||||
template <typename T, typename FUNC>
|
||||
int64_t
|
||||
ProcessAllDataChunk(
|
||||
FUNC func,
|
||||
std::function<bool(const milvus::SkipIndex&, FieldId, int)> skip_func,
|
||||
TargetBitmapView res,
|
||||
TargetBitmapView valid_res,
|
||||
const ValTypes&... values) {
|
||||
if (segment_->is_chunked()) {
|
||||
return ProcessAllChunksForMultipleChunk<T>(
|
||||
func, skip_func, res, valid_res, values...);
|
||||
} else {
|
||||
ThrowInfo(ErrorCode::Unsupported, "unreachable");
|
||||
ProcessAllDataChunkBatched(FUNC func, TargetBitmapView res) {
|
||||
static_assert(std::is_same_v<T, std::string_view> ||
|
||||
std::is_same_v<T, Json> ||
|
||||
std::is_same_v<T, ArrayView>,
|
||||
"ProcessAllDataChunkBatched only supports string_view, "
|
||||
"Json, and ArrayView types");
|
||||
|
||||
AssertInfo(segment_->is_chunked(),
|
||||
"ProcessAllDataChunkBatched requires chunked segment");
|
||||
AssertInfo(segment_->type() == SegmentType::Sealed,
|
||||
"ProcessAllDataChunkBatched requires sealed segment");
|
||||
|
||||
int64_t processed_size = 0;
|
||||
|
||||
for (size_t chunk_id = 0; chunk_id < num_data_chunk_; chunk_id++) {
|
||||
int64_t chunk_size = segment_->chunk_size(field_id_, chunk_id);
|
||||
int64_t chunk_offset = 0;
|
||||
|
||||
while (chunk_offset < chunk_size) {
|
||||
int64_t batch_size =
|
||||
std::min(batch_size_, chunk_size - chunk_offset);
|
||||
|
||||
auto pw = segment_->get_batch_views<T>(
|
||||
op_ctx_, field_id_, chunk_id, chunk_offset, batch_size);
|
||||
auto data_vec = std::move(pw.get().first);
|
||||
|
||||
func(data_vec.data(), batch_size, res + processed_size);
|
||||
|
||||
chunk_offset += batch_size;
|
||||
processed_size += batch_size;
|
||||
}
|
||||
}
|
||||
|
||||
return processed_size;
|
||||
}
|
||||
|
||||
int
|
||||
|
||||
@@ -289,16 +289,16 @@ NgramInvertedIndex::ExecuteQuery(const std::string& literal,
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
template <typename T, typename Predicate>
|
||||
inline void
|
||||
handle_batch(const T* data,
|
||||
const int size,
|
||||
const int64_t size,
|
||||
TargetBitmapView res,
|
||||
std::function<bool(const T&)> predicate) {
|
||||
Predicate&& predicate) {
|
||||
auto next_off_option = res.find_first();
|
||||
while (next_off_option.has_value()) {
|
||||
auto next_off = next_off_option.value();
|
||||
if (next_off >= size) {
|
||||
if (next_off >= static_cast<size_t>(size)) {
|
||||
return;
|
||||
}
|
||||
if (!predicate(data[next_off])) {
|
||||
@@ -308,13 +308,12 @@ handle_batch(const T* data,
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
template <typename T, typename Predicate>
|
||||
std::optional<TargetBitmap>
|
||||
NgramInvertedIndex::ExecuteQueryWithPredicate(
|
||||
const std::string& literal,
|
||||
exec::SegmentExpr* segment,
|
||||
std::function<bool(const T&)> predicate,
|
||||
bool need_post_filter) {
|
||||
NgramInvertedIndex::ExecuteQueryWithPredicate(const std::string& literal,
|
||||
exec::SegmentExpr* segment,
|
||||
Predicate&& predicate,
|
||||
bool need_post_filter) {
|
||||
TargetBitmap bitset{static_cast<size_t>(Count())};
|
||||
wrapper_->ngram_match_query(literal, min_gram_, max_gram_, &bitset);
|
||||
|
||||
@@ -323,24 +322,15 @@ NgramInvertedIndex::ExecuteQueryWithPredicate(
|
||||
|
||||
if (need_post_filter) {
|
||||
TargetBitmapView res(bitset);
|
||||
TargetBitmap valid(res.size(), true);
|
||||
TargetBitmapView valid_res(valid.data(), valid.size());
|
||||
|
||||
auto execute_batch =
|
||||
[&predicate](
|
||||
const T* data,
|
||||
// `valid_data` is not used as the results returned by ngram_match_query are all valid
|
||||
const bool* _valid_data,
|
||||
const int32_t* _offsets,
|
||||
const int size,
|
||||
TargetBitmapView res,
|
||||
// the same with `valid_data`
|
||||
TargetBitmapView _valid_res) {
|
||||
handle_batch(data, size, res, predicate);
|
||||
};
|
||||
auto execute_batch = [&predicate](const T* data,
|
||||
const int64_t size,
|
||||
TargetBitmapView res) {
|
||||
handle_batch(data, size, res, predicate);
|
||||
};
|
||||
|
||||
segment->template ProcessAllDataChunkBatched<T>(execute_batch, res);
|
||||
|
||||
segment->ProcessAllDataChunk<T>(
|
||||
execute_batch, std::nullptr_t{}, res, valid_res);
|
||||
final_result_count = bitset.count();
|
||||
}
|
||||
|
||||
@@ -407,8 +397,6 @@ NgramInvertedIndex::MatchQuery(const std::string& literal,
|
||||
}
|
||||
|
||||
TargetBitmapView res(bitset);
|
||||
TargetBitmap valid(res.size(), true);
|
||||
TargetBitmapView valid_res(valid.data(), valid.size());
|
||||
|
||||
PatternMatchTranslator translator;
|
||||
auto regex_pattern = translator(literal);
|
||||
@@ -417,50 +405,35 @@ NgramInvertedIndex::MatchQuery(const std::string& literal,
|
||||
auto ngram_hit_count = bitset.count();
|
||||
|
||||
if (schema_.data_type() == proto::schema::DataType::JSON) {
|
||||
auto predicate = [&literal, &matcher, this](const milvus::Json& data) {
|
||||
auto predicate = [&matcher, this](const milvus::Json& data) {
|
||||
auto x = data.template at<std::string_view>(this->nested_path_);
|
||||
if (x.error()) {
|
||||
return false;
|
||||
}
|
||||
auto data_val = x.value();
|
||||
return matcher(data_val);
|
||||
return matcher(x.value());
|
||||
};
|
||||
|
||||
auto execute_batch_json =
|
||||
[&predicate](
|
||||
const milvus::Json* data,
|
||||
// `valid_data` is not used as the results returned by ngram_match_query are all valid
|
||||
const bool* _valid_data,
|
||||
const int32_t* _offsets,
|
||||
const int size,
|
||||
TargetBitmapView res,
|
||||
// the same with `valid_data`
|
||||
TargetBitmapView _valid_res,
|
||||
std::string val) {
|
||||
handle_batch<milvus::Json>(data, size, res, predicate);
|
||||
};
|
||||
auto execute_batch = [&predicate](const milvus::Json* data,
|
||||
const int64_t size,
|
||||
TargetBitmapView res) {
|
||||
handle_batch<milvus::Json>(data, size, res, predicate);
|
||||
};
|
||||
|
||||
segment->ProcessAllDataChunk<milvus::Json>(
|
||||
execute_batch_json, std::nullptr_t{}, res, valid_res, literal);
|
||||
segment->template ProcessAllDataChunkBatched<milvus::Json>(
|
||||
execute_batch, res);
|
||||
} else {
|
||||
auto predicate = [&matcher](const std::string_view& data) {
|
||||
return matcher(data);
|
||||
};
|
||||
|
||||
auto execute_batch =
|
||||
[&predicate](
|
||||
const std::string_view* data,
|
||||
// `valid_data` is not used as the results returned by ngram_match_query are all valid
|
||||
const bool* _valid_data,
|
||||
const int32_t* _offsets,
|
||||
const int size,
|
||||
TargetBitmapView res,
|
||||
// the same with `valid_data`
|
||||
TargetBitmapView _valid_res) {
|
||||
handle_batch<std::string_view>(data, size, res, predicate);
|
||||
};
|
||||
segment->ProcessAllDataChunk<std::string_view>(
|
||||
execute_batch, std::nullptr_t{}, res, valid_res);
|
||||
auto execute_batch = [&predicate](const std::string_view* data,
|
||||
const int64_t size,
|
||||
TargetBitmapView res) {
|
||||
handle_batch<std::string_view>(data, size, res, predicate);
|
||||
};
|
||||
|
||||
segment->template ProcessAllDataChunkBatched<std::string_view>(
|
||||
execute_batch, res);
|
||||
}
|
||||
|
||||
auto final_result_count = bitset.count();
|
||||
|
||||
@@ -65,11 +65,11 @@ class NgramInvertedIndex : public InvertedIndexTantivy<std::string> {
|
||||
}
|
||||
|
||||
private:
|
||||
template <typename T>
|
||||
template <typename T, typename Predicate>
|
||||
std::optional<TargetBitmap>
|
||||
ExecuteQueryWithPredicate(const std::string& literal,
|
||||
exec::SegmentExpr* segment,
|
||||
std::function<bool(const T&)> predicate,
|
||||
Predicate&& predicate,
|
||||
bool need_post_filter);
|
||||
|
||||
// Match is something like xxx%xxx%xxx, xxx%xxx, %xxx%xxx, xxx_x etc.
|
||||
|
||||
Reference in New Issue
Block a user