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:
Spade A
2025-12-30 16:57:22 +08:00
committed by GitHub
parent a1721bb47b
commit 1a6f3c4305
3 changed files with 79 additions and 117 deletions
+44 -55
View File
@@ -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
+33 -60
View File
@@ -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();
+2 -2
View File
@@ -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.