Skip to content

Commit d1debd8

Browse files
hicder-rddthicder-rddt
committed
RDDT: Group-by loop on proxy
1 parent 07aa2c4 commit d1debd8

24 files changed

Lines changed: 1441 additions & 93 deletions

‎go.sum‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -984,6 +984,7 @@ github.com/shoenig/go-m1cpu v0.2.2 h1:4nc55oVv7nygGnfI9bhLCLzUEs4794y0Bkqx4q2zy7
984984
github.com/shoenig/go-m1cpu v0.2.2/go.mod h1:KkDOw6m3ZJQAPHbrzkZki4hnx+pDRR1Lo+ldA56wD5w=
985985
github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU=
986986
github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
987+
github.com/shoenig/test v1.7.0 h1:eWcHtTXa6QLnBvm0jgEabMRN/uJ4DMV3M8xUGgRkZmk=
987988
github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME=
988989
github.com/shurcooL/go v0.0.0-20180423040247-9e1955d9fb6e/go.mod h1:TDJrrUr11Vxrven61rcy3hJMUqaf/CLWYhHNPmT14Lk=
989990
github.com/shurcooL/httpfs v0.0.0-20181222201310-74dc9339e414/go.mod h1:ZY1cvUeJuFPAdZ/B6v7RHavJWZn2YPVFQ1OSXhCGOkg=

‎internal/core/src/common/Consts.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ const char VEC_OPT_FIELDS[] = "opt_fields";
5151
const char PAGE_RETAIN_ORDER[] = "page_retain_order";
5252
const char TEXT_LOG_ROOT_PATH[] = "text_log";
5353
const char ITERATIVE_FILTER[] = "iterative_filter";
54+
const char GROUP_BY_REFILL[] = "group_by_refill";
5455
const char HINTS[] = "hints";
5556
// json stats related
5657
const char JSON_KEY_INDEX_LOG_ROOT_PATH[] = "json_key_index_log";

‎internal/core/src/common/QueryInfo.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ struct SearchInfo {
4343
tracer::TraceContext trace_ctx_;
4444
bool materialized_view_involved = false;
4545
bool iterative_filter_execution = false;
46+
bool proxy_group_by_refill_ = false;
4647
std::optional<SearchIteratorV2Info> iterator_v2_info_ = std::nullopt;
4748
std::optional<std::string> json_path_;
4849
std::optional<milvus::DataType> json_type_;

‎internal/core/src/exec/operator/GroupByNode.cpp‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,14 @@
1515
// limitations under the License.
1616

1717
#include "GroupByNode.h"
18+
#include <numeric>
19+
1820
#include "common/Tracer.h"
1921
#include "fmt/format.h"
2022

23+
#include "exec/operator/Utils.h"
2124
#include "exec/operator/groupby/SearchGroupByOperator.h"
25+
#include "log/Log.h"
2226
#include "monitor/Monitor.h"
2327
namespace milvus {
2428
namespace exec {
@@ -66,7 +70,8 @@ PhyGroupByNode::GetOutput() {
6670

6771
auto op_context = query_context_->get_op_context();
6872
auto search_result = query_context_->get_search_result();
69-
if (search_result.vector_iterators_.has_value()) {
73+
if (UseVectorIterator(search_info_) &&
74+
search_result.vector_iterators_.has_value()) {
7075
AssertInfo(search_result.vector_iterators_.value().size() ==
7176
search_result.total_nq_,
7277
"Vector Iterators' count must be equal to total_nq_, Check "
@@ -88,6 +93,25 @@ PhyGroupByNode::GetOutput() {
8893
"equal to search_result.seg_offsets.size:{}",
8994
search_result.group_by_values_.value().size(),
9095
search_result.seg_offsets_.size());
96+
} else {
97+
std::vector<GroupByValueType> group_by_values;
98+
milvus::exec::PopulateGroupByValues(op_context,
99+
search_info_,
100+
group_by_values,
101+
*segment_,
102+
search_result.seg_offsets_);
103+
search_result.group_by_values_ = std::move(group_by_values);
104+
search_result.group_size_ = search_info_.group_size_;
105+
if (search_result.topk_per_nq_prefix_sum_.empty()) {
106+
std::vector<size_t> topks(search_result.total_nq_,
107+
search_result.unity_topK_);
108+
search_result.topk_per_nq_prefix_sum_.resize(
109+
search_result.total_nq_ + 1);
110+
std::partial_sum(
111+
topks.begin(),
112+
topks.end(),
113+
search_result.topk_per_nq_prefix_sum_.begin() + 1);
114+
}
91115
}
92116
tracer::AddEvent(
93117
fmt::format("grouped_results: {}", search_result.seg_offsets_.size()));
@@ -109,4 +133,4 @@ PhyGroupByNode::IsFinished() {
109133
}
110134

111135
} // namespace exec
112-
} // namespace milvus
136+
} // namespace milvus

‎internal/core/src/exec/operator/Utils.h‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,16 @@ namespace exec {
3333

3434
static bool
3535
UseVectorIterator(const SearchInfo& search_info) {
36-
return search_info.group_by_field_id_.has_value() ||
37-
search_info.iterative_filter_execution;
36+
if (search_info.iterative_filter_execution) {
37+
return true;
38+
}
39+
// Default group-by uses vector iterators at core. Proxy-side refill disables
40+
// this path and handles incomplete groups with follow-up searches instead.
41+
if (search_info.group_by_field_id_.has_value() &&
42+
!search_info.proxy_group_by_refill_) {
43+
return true;
44+
}
45+
return false;
3846
}
3947

4048
static bool

‎internal/core/src/exec/operator/groupby/SearchGroupByOperator.cpp‎

Lines changed: 182 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,187 @@ SearchGroupBy(milvus::OpContext* op_ctx,
316316
}
317317
}
318318

319+
template <typename T>
320+
static void
321+
PopulateGroupByValuesByType(const std::shared_ptr<DataGetter<T>>& data_getter,
322+
std::vector<GroupByValueType>& group_by_values,
323+
const std::vector<int64_t>& seg_offsets) {
324+
group_by_values.reserve(seg_offsets.size());
325+
for (const auto offset : seg_offsets) {
326+
if (offset == INVALID_SEG_OFFSET) {
327+
group_by_values.emplace_back(std::nullopt);
328+
continue;
329+
}
330+
group_by_values.emplace_back(data_getter->Get(offset));
331+
}
332+
}
333+
334+
void
335+
PopulateGroupByValues(milvus::OpContext* op_ctx,
336+
const SearchInfo& search_info,
337+
std::vector<GroupByValueType>& group_by_values,
338+
const segcore::SegmentInternalInterface& segment,
339+
const std::vector<int64_t>& seg_offsets) {
340+
FieldId group_by_field_id = search_info.group_by_field_id_.value();
341+
auto data_type = segment.GetFieldDataType(group_by_field_id);
342+
switch (data_type) {
343+
case DataType::INT8: {
344+
PopulateGroupByValuesByType<int8_t>(
345+
GetDataGetter<int8_t>(op_ctx, segment, group_by_field_id),
346+
group_by_values,
347+
seg_offsets);
348+
break;
349+
}
350+
case DataType::INT16: {
351+
PopulateGroupByValuesByType<int16_t>(
352+
GetDataGetter<int16_t>(op_ctx, segment, group_by_field_id),
353+
group_by_values,
354+
seg_offsets);
355+
break;
356+
}
357+
case DataType::INT32: {
358+
PopulateGroupByValuesByType<int32_t>(
359+
GetDataGetter<int32_t>(op_ctx, segment, group_by_field_id),
360+
group_by_values,
361+
seg_offsets);
362+
break;
363+
}
364+
case DataType::INT64:
365+
case DataType::TIMESTAMPTZ: {
366+
PopulateGroupByValuesByType<int64_t>(
367+
GetDataGetter<int64_t>(op_ctx, segment, group_by_field_id),
368+
group_by_values,
369+
seg_offsets);
370+
break;
371+
}
372+
case DataType::BOOL: {
373+
PopulateGroupByValuesByType<bool>(
374+
GetDataGetter<bool>(op_ctx, segment, group_by_field_id),
375+
group_by_values,
376+
seg_offsets);
377+
break;
378+
}
379+
case DataType::VARCHAR: {
380+
PopulateGroupByValuesByType<std::string>(
381+
GetDataGetter<std::string>(op_ctx, segment, group_by_field_id),
382+
group_by_values,
383+
seg_offsets);
384+
break;
385+
}
386+
case DataType::JSON: {
387+
AssertInfo(search_info.json_path_.has_value(),
388+
"json_path is required for json field when doing "
389+
"search_group_by");
390+
if (search_info.json_type_.has_value()) {
391+
switch (search_info.json_type_.value()) {
392+
case DataType::BOOL: {
393+
PopulateGroupByValuesByType<bool>(
394+
GetDataGetter<bool, milvus::Json>(
395+
op_ctx,
396+
segment,
397+
group_by_field_id,
398+
search_info.json_path_,
399+
search_info.json_type_,
400+
search_info.strict_cast_),
401+
group_by_values,
402+
seg_offsets);
403+
break;
404+
}
405+
case DataType::INT8: {
406+
PopulateGroupByValuesByType<int8_t>(
407+
GetDataGetter<int8_t, milvus::Json>(
408+
op_ctx,
409+
segment,
410+
group_by_field_id,
411+
search_info.json_path_,
412+
search_info.json_type_,
413+
search_info.strict_cast_),
414+
group_by_values,
415+
seg_offsets);
416+
break;
417+
}
418+
case DataType::INT16: {
419+
PopulateGroupByValuesByType<int16_t>(
420+
GetDataGetter<int16_t, milvus::Json>(
421+
op_ctx,
422+
segment,
423+
group_by_field_id,
424+
search_info.json_path_,
425+
search_info.json_type_,
426+
search_info.strict_cast_),
427+
group_by_values,
428+
seg_offsets);
429+
break;
430+
}
431+
case DataType::INT32: {
432+
PopulateGroupByValuesByType<int32_t>(
433+
GetDataGetter<int32_t, milvus::Json>(
434+
op_ctx,
435+
segment,
436+
group_by_field_id,
437+
search_info.json_path_,
438+
search_info.json_type_,
439+
search_info.strict_cast_),
440+
group_by_values,
441+
seg_offsets);
442+
break;
443+
}
444+
case DataType::INT64: {
445+
PopulateGroupByValuesByType<int64_t>(
446+
GetDataGetter<int64_t, milvus::Json>(
447+
op_ctx,
448+
segment,
449+
group_by_field_id,
450+
search_info.json_path_,
451+
search_info.json_type_,
452+
search_info.strict_cast_),
453+
group_by_values,
454+
seg_offsets);
455+
break;
456+
}
457+
case DataType::VARCHAR: {
458+
PopulateGroupByValuesByType<std::string>(
459+
GetDataGetter<std::string, milvus::Json>(
460+
op_ctx,
461+
segment,
462+
group_by_field_id,
463+
search_info.json_path_,
464+
search_info.json_type_,
465+
search_info.strict_cast_),
466+
group_by_values,
467+
seg_offsets);
468+
break;
469+
}
470+
default: {
471+
ThrowInfo(Unsupported,
472+
fmt::format("unsupported data type {} for "
473+
"group by operator",
474+
data_type));
475+
}
476+
}
477+
} else {
478+
PopulateGroupByValuesByType<std::string>(
479+
GetDataGetter<std::string, milvus::Json>(
480+
op_ctx,
481+
segment,
482+
group_by_field_id,
483+
search_info.json_path_,
484+
search_info.json_type_,
485+
search_info.strict_cast_),
486+
group_by_values,
487+
seg_offsets);
488+
}
489+
break;
490+
}
491+
default: {
492+
ThrowInfo(
493+
Unsupported,
494+
fmt::format("unsupported data type {} for group by operator",
495+
data_type));
496+
}
497+
}
498+
}
499+
319500
template <typename T>
320501
void
321502
GroupIteratorsByType(
@@ -392,4 +573,4 @@ GroupIteratorResult(const std::shared_ptr<VectorIterator>& iterator,
392573
}
393574

394575
} // namespace exec
395-
} // namespace milvus
576+
} // namespace milvus

‎internal/core/src/exec/operator/groupby/SearchGroupByOperator.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,13 @@ SearchGroupBy(milvus::OpContext* op_ctx,
325325
std::vector<float>& distances,
326326
std::vector<size_t>& topk_per_nq_prefix_sum);
327327

328+
void
329+
PopulateGroupByValues(milvus::OpContext* op_ctx,
330+
const SearchInfo& searchInfo,
331+
std::vector<GroupByValueType>& group_by_values,
332+
const segcore::SegmentInternalInterface& segment,
333+
const std::vector<int64_t>& seg_offsets);
334+
328335
template <typename T>
329336
void
330337
GroupIteratorsByType(

‎internal/core/src/query/PlanProto.cpp‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,29 @@ ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) {
9191
}
9292
}
9393

94+
if (search_info.search_params_.contains(GROUP_BY_REFILL)) {
95+
auto& refill_param = search_info.search_params_[GROUP_BY_REFILL];
96+
if (refill_param.is_boolean()) {
97+
search_info.proxy_group_by_refill_ = refill_param.get<bool>();
98+
} else if (refill_param.is_string()) {
99+
auto refill_param_str = refill_param.get<std::string>();
100+
if (refill_param_str == "true" || refill_param_str == "True") {
101+
search_info.proxy_group_by_refill_ = true;
102+
} else if (refill_param_str == "false" ||
103+
refill_param_str == "False") {
104+
search_info.proxy_group_by_refill_ = false;
105+
} else {
106+
ThrowInfo(ConfigInvalid,
107+
"group_by_refill: {} not supported",
108+
refill_param);
109+
}
110+
} else {
111+
ThrowInfo(ConfigInvalid,
112+
"group_by_refill: {} not supported",
113+
refill_param);
114+
}
115+
}
116+
94117
if (query_info_proto.bm25_avgdl() > 0) {
95118
search_info.search_params_[knowhere::meta::BM25_AVGDL] =
96119
query_info_proto.bm25_avgdl();

‎internal/core/src/segcore/ReduceUtils.cpp‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,13 @@ AssembleGroupByValues(
138138
mutable_group_by_field_value->set_type(
139139
milvus::proto::schema::DataType::VarChar);
140140
auto field_data = group_by_res_values->mutable_string_data();
141+
auto mutable_data = field_data->mutable_data();
142+
for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
143+
mutable_data->Add();
144+
}
141145
for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
142146
if (group_by_vals[idx].has_value()) {
143-
*(field_data->mutable_data()->Add()) =
147+
*mutable_data->Mutable(idx) =
144148
std::get<std::string>(group_by_vals[idx].value());
145149
} else {
146150
valid_data->Set(idx, false);
@@ -152,9 +156,13 @@ AssembleGroupByValues(
152156
mutable_group_by_field_value->set_type(
153157
milvus::proto::schema::DataType::Geometry);
154158
auto field_data = group_by_res_values->mutable_geometry_data();
159+
auto mutable_data = field_data->mutable_data();
160+
for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
161+
mutable_data->Add();
162+
}
155163
for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
156164
if (group_by_vals[idx].has_value()) {
157-
*(field_data->mutable_data()->Add()) =
165+
*mutable_data->Mutable(idx) =
158166
std::get<std::string>(group_by_vals[idx].value());
159167
} else {
160168
valid_data->Set(idx, false);
@@ -269,10 +277,15 @@ AssembleGroupByValues(
269277
milvus::proto::schema::DataType::VarChar);
270278
auto field_data =
271279
group_by_res_values->mutable_string_data();
280+
auto mutable_data = field_data->mutable_data();
281+
for (std::size_t idx = 0; idx < group_by_val_size;
282+
idx++) {
283+
mutable_data->Add();
284+
}
272285
for (std::size_t idx = 0; idx < group_by_val_size;
273286
idx++) {
274287
if (group_by_vals[idx].has_value()) {
275-
*(field_data->mutable_data()->Add()) =
288+
*mutable_data->Mutable(idx) =
276289
std::get<std::string>(
277290
group_by_vals[idx].value());
278291
} else {
@@ -305,4 +318,4 @@ AssembleGroupByValues(
305318
}
306319
}
307320

308-
} // namespace milvus::segcore
321+
} // namespace milvus::segcore

0 commit comments

Comments
 (0)