Skip to content

[feature](query) Generate TopN filters from aggregation #66652

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
HappenLee wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[feature](query) Generate TopN filters from aggregation #66652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions be/src/exec/operator/aggregation_sink_operator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "exprs/aggregate/aggregate_function_count.h"
#include "exprs/aggregate/aggregate_function_simple_factory.h"
#include "exprs/vectorized_agg_fn.h"
#include "runtime/query_context.h"
#include "runtime/runtime_profile.h"
#include "runtime/thread_context.h"

Expand Down Expand Up @@ -77,6 +78,7 @@ Status AggSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& info) {

_memory_usage_container = ADD_COUNTER(custom_profile(), "MemoryUsageContainer", TUnit::BYTES);
_memory_usage_arena = ADD_COUNTER(custom_profile(), "MemoryUsageArena", TUnit::BYTES);
_update_runtime_predicate_timer = ADD_TIMER(custom_profile(), "UpdateRuntimePredicateTime");

return Status::OK();
}
Expand Down Expand Up @@ -225,6 +227,31 @@ Status AggSinkLocalState::_merge_with_serialized_key(Block* block) {
}
}

Status AggSinkLocalState::_update_runtime_predicate(RuntimeState* state) {
auto& p = Base::_parent->template cast<AggSinkOperatorX>();
auto* query_ctx = state->get_query_ctx();
if (!query_ctx->has_runtime_predicate(p.node_id()) || !Base::_shared_state->reach_limit) {
return Status::OK();
}

SCOPED_TIMER(_update_runtime_predicate_timer);
auto& predicate = query_ctx->get_runtime_predicate(p.node_id());
if (!predicate.enable()) {
return Status::OK();
}

DCHECK(Base::_shared_state->do_sort_limit);
DCHECK(!Base::_shared_state->limit_columns.empty());
DCHECK_GE(Base::_shared_state->limit_columns_min, 0);
Field new_top {PrimitiveType::TYPE_NULL};
Base::_shared_state->limit_columns[0]->get(Base::_shared_state->limit_columns_min, new_top);
if (!new_top.is_null() && new_top != _old_top) {
RETURN_IF_ERROR(predicate.update(new_top));
_old_top = std::move(new_top);
}
return Status::OK();
}

size_t AggSinkLocalState::_memory_usage() const {
if (0 == get_hash_table_size()) {
return 0;
Expand Down Expand Up @@ -895,6 +922,13 @@ Status AggSinkOperatorX::init(const TPlanNode& tnode, RuntimeState* state) {
}
}

auto* query_ctx = state->get_query_ctx();
if (query_ctx->has_runtime_predicate(_node_id)) {
DORIS_CHECK(_do_sort_limit);
DORIS_CHECK_GT(_limit, 0);
query_ctx->get_runtime_predicate(_node_id).set_detected_source();
}

return Status::OK();
}

Expand Down Expand Up @@ -991,6 +1025,7 @@ Status AggSinkOperatorX::sink_impl(doris::RuntimeState* state, Block* in_block,
local_state._shared_state->input_num_rows += in_block->rows();
if (in_block->rows() > 0) {
RETURN_IF_ERROR(local_state._executor->execute(&local_state, in_block));
RETURN_IF_ERROR(local_state._update_runtime_predicate(state));
local_state._executor->update_memusage(&local_state);
COUNTER_SET(local_state._hash_table_size_counter,
(int64_t)local_state.get_hash_table_size());
Expand Down
4 changes: 4 additions & 0 deletions be/src/exec/operator/aggregation_sink_operator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

#include <stdint.h>

#include "core/field.h"
#include "exec/operator/operator.h"
#include "runtime/exec_env.h"
#include "runtime/runtime_profile.h"
Expand Down Expand Up @@ -80,6 +81,7 @@ class AggSinkLocalState : public PipelineXSinkLocalState<AggSharedState> {
Status _execute_with_serialized_key(Block* block);
Status _merge_with_serialized_key(Block* block);
void _update_memusage_with_serialized_key();
Status _update_runtime_predicate(RuntimeState* state);
template <bool limit>

Status _execute_with_serialized_key_helper(Block* block);
Expand Down Expand Up @@ -117,6 +119,7 @@ class AggSinkLocalState : public PipelineXSinkLocalState<AggSharedState> {
RuntimeProfile::Counter* _serialize_key_arena_memory_usage = nullptr;
RuntimeProfile::Counter* _memory_usage_container = nullptr;
RuntimeProfile::Counter* _memory_usage_arena = nullptr;
RuntimeProfile::Counter* _update_runtime_predicate_timer = nullptr;

bool _should_limit_output = false;

Expand All @@ -130,6 +133,7 @@ class AggSinkLocalState : public PipelineXSinkLocalState<AggSharedState> {
std::unique_ptr<ExecutorBase> _executor = nullptr;

int64_t _memory_usage_last_executing = 0;
Field _old_top {PrimitiveType::TYPE_NULL};
};

class AggSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX<AggSinkLocalState> {
Expand Down
35 changes: 35 additions & 0 deletions be/src/exec/operator/streaming_aggregation_operator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
#include "exprs/aggregate/aggregate_function_simple_factory.h"
#include "exprs/vectorized_agg_fn.h"
#include "exprs/vslot_ref.h"
#include "runtime/query_context.h"

namespace doris {
class RuntimeState;
Expand Down Expand Up @@ -72,6 +73,7 @@ Status StreamingAggLocalState::init(RuntimeState* state, LocalStateInfo& info) {
_get_results_timer = ADD_TIMER(custom_profile(), "GetResultsTime");
_hash_table_iterate_timer = ADD_TIMER(custom_profile(), "HashTableIterateTime");
_insert_keys_to_column_timer = ADD_TIMER(custom_profile(), "InsertKeysToColumnTime");
_update_runtime_predicate_timer = ADD_TIMER(custom_profile(), "UpdateRuntimePredicateTime");

return Status::OK();
}
Expand Down Expand Up @@ -192,6 +194,31 @@ Status StreamingAggLocalState::do_pre_agg(RuntimeState* state, Block* input_bloc
return Status::OK();
}

Status StreamingAggLocalState::_update_runtime_predicate(RuntimeState* state) {
auto& p = Base::_parent->template cast<StreamingAggOperatorX>();
auto* query_ctx = state->get_query_ctx();
if (!query_ctx->has_runtime_predicate(p.node_id()) || need_do_sort_limit != 1) {
return Status::OK();
}

SCOPED_TIMER(_update_runtime_predicate_timer);
auto& predicate = query_ctx->get_runtime_predicate(p.node_id());
if (!predicate.enable()) {
return Status::OK();
}

DCHECK(do_sort_limit);
DCHECK(!limit_columns.empty());
DCHECK_GE(limit_columns_min, 0);
Field new_top {PrimitiveType::TYPE_NULL};
limit_columns[0]->get(limit_columns_min, new_top);
if (!new_top.is_null() && new_top != _old_top) {

@github-actions github-actions Bot Aug 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Do not publish this boundary after passthrough has allowed duplicate group keys into limit_heap. The heap is initially built from distinct hash-table keys, but _add_limit_heap_top later inserts the first better raw passthrough row without checking whether that key is already represented. For LIMIT 3 ASC, {1,2,100} followed by another row for key 1 becomes {1,1,2}; publishing <=2 through the shared QueryContext can make another scan task drop key 50, although distinct groups {1,2,50,100} still need 50 in the top three. Please keep publication backed by LIMIT distinct keys (or stop tightening after passthrough begins), and cover a duplicate-better-key plus cross-instance case.

RETURN_IF_ERROR(predicate.update(new_top));
_old_top = std::move(new_top);
}
return Status::OK();
}

bool StreamingAggLocalState::_should_expand_preagg_hash_tables() {
if (!_should_expand_hash_table) {
return false;
Expand Down Expand Up @@ -992,6 +1019,13 @@ Status StreamingAggOperatorX::init(const TPlanNode& tnode, RuntimeState* state)
}
}

auto* query_ctx = state->get_query_ctx();
if (query_ctx->has_runtime_predicate(_node_id)) {
DORIS_CHECK(_do_sort_limit);
DORIS_CHECK_GT(_sort_limit, 0);
query_ctx->get_runtime_predicate(_node_id).set_detected_source();
}

_op_name = "STREAMING_AGGREGATION_OPERATOR";
return Status::OK();
}
Expand Down Expand Up @@ -1124,6 +1158,7 @@ Status StreamingAggOperatorX::push(RuntimeState* state, Block* in_block, bool eo
if (in_block->rows() > 0) {
RETURN_IF_ERROR(
local_state.do_pre_agg(state, in_block, local_state._pre_aggregated_block.get()));
RETURN_IF_ERROR(local_state._update_runtime_predicate(state));
}
in_block->clear_column_data(_child->row_desc().num_materialized_slots());
return Status::OK();
Expand Down
4 changes: 4 additions & 0 deletions be/src/exec/operator/streaming_aggregation_operator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

#include "common/status.h"
#include "core/block/block.h"
#include "core/field.h"
#include "exec/operator/operator.h"
#include "runtime/runtime_profile.h"

Expand Down Expand Up @@ -55,6 +56,7 @@ class StreamingAggLocalState MOCK_REMOVE(final) : public PipelineXLocalState<Fak
void _add_limit_heap_top(ColumnRawPtrs& key_columns, size_t rows);
bool _do_limit_filter(size_t num_rows, ColumnRawPtrs& key_columns);
void _refresh_limit_heap(size_t i, ColumnRawPtrs& key_columns);
Status _update_runtime_predicate(RuntimeState* state);

Status _pre_agg_with_serialized_key(doris::Block* in_block, doris::Block* out_block);
bool _should_expand_preagg_hash_tables();
Expand Down Expand Up @@ -89,6 +91,7 @@ class StreamingAggLocalState MOCK_REMOVE(final) : public PipelineXLocalState<Fak
RuntimeProfile::Counter* _get_results_timer = nullptr;
RuntimeProfile::Counter* _hash_table_iterate_timer = nullptr;
RuntimeProfile::Counter* _insert_keys_to_column_timer = nullptr;
RuntimeProfile::Counter* _update_runtime_predicate_timer = nullptr;

bool _should_expand_hash_table = true;
int64_t _cur_num_rows_returned = 0;
Expand All @@ -107,6 +110,7 @@ class StreamingAggLocalState MOCK_REMOVE(final) : public PipelineXLocalState<Fak
bool do_sort_limit = false;
MutableColumns limit_columns;
int limit_columns_min = -1;
Field _old_top {PrimitiveType::TYPE_NULL};
PaddedPODArray<uint8_t> need_computes;
std::vector<uint8_t> cmp_res;
std::vector<int> order_directions;
Expand Down
118 changes: 117 additions & 1 deletion be/test/exec/operator/agg_operator_group_by_limit_opt_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,16 +27,43 @@
#include "exec/operator/assert_num_rows_operator.h"
#include "exec/operator/mock_operator.h"
#include "exec/operator/operator_helper.h"
#include "exec/pipeline/thrift_builder.h"
#include "runtime/query_context.h"
#include "testutil/column_helper.h"
#include "testutil/mock/mock_agg_fn_evaluator.h"
#include "testutil/mock/mock_slot_ref.h"

namespace doris {
namespace {

constexpr TPlanNodeId TOPN_FILTER_TARGET_NODE_ID = 20;

void init_topn_runtime_predicate(MockRuntimeState& state, TPlanNodeId source_node_id, bool is_asc) {
auto target_expr = TRuntimeFilterDescBuilder::get_default_expr();
target_expr.nodes[0].__set_type(create_type_desc(TYPE_BIGINT));

TTopnFilterDesc desc;
desc.__set_source_node_id(source_node_id);
desc.__set_is_asc(is_asc);
desc.__set_null_first(false);
desc.__set_target_node_id_to_target_expr({{TOPN_FILTER_TARGET_NODE_ID, target_expr}});
state.get_query_ctx()->init_runtime_predicates({desc});

auto& predicate = state.get_query_ctx()->get_runtime_predicate(source_node_id);
predicate.set_detected_source();
DORIS_CHECK(predicate.init_target(TOPN_FILTER_TARGET_NODE_ID, {}, -1).ok());
}

Field get_topn_runtime_predicate_value(MockRuntimeState& state, TPlanNodeId source_node_id) {
return state.get_query_ctx()->get_runtime_predicate(source_node_id).get_value();
}

} // namespace

auto static init_sink_and_source(std::shared_ptr<AggSinkOperatorX> sink_op,
std::shared_ptr<AggSourceOperatorX> source_op,
OperatorContext& ctx) {
auto shared_state = sink_op->create_shared_state();
auto shared_state = std::static_pointer_cast<AggSharedState>(sink_op->create_shared_state());
{
auto local_state = AggSinkOperatorX::LocalState ::create_unique(sink_op.get(), &ctx.state);
LocalSinkStateInfo info {.task_idx = 0,
Expand Down Expand Up @@ -194,6 +221,90 @@ TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_need_finalize_without_ord
}
}

TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_desc_order_by_updates_runtime_predicate) {
OperatorContext ctx;
auto sink_op = std::make_shared<MockAggsinkOperator>();
sink_op->_aggregate_evaluators.push_back(create_mock_agg_fn_evaluator(
ctx.pool, MockSlotRef::create_mock_contexts(1, std::make_shared<DataTypeInt64>()),
false, false));
sink_op->_pool = &ctx.pool;
sink_op->_limit = 2;
sink_op->_probe_expr_ctxs =
MockSlotRef::create_mock_contexts(0, std::make_shared<DataTypeInt64>());
sink_op->_order_directions = {-1};
sink_op->_null_directions = {-1};
sink_op->_do_sort_limit = true;
ASSERT_TRUE(sink_op->prepare(&ctx.state).ok());

auto source_op = std::make_shared<MockAggSourceOperator>();
source_op->mock_row_descriptor.reset(new MockRowDescriptor {
{std::make_shared<DataTypeInt64>(), std::make_shared<DataTypeInt64>()}, &ctx.pool});
source_op->_without_key = false;
source_op->_needs_finalize = true;
ASSERT_TRUE(source_op->prepare(&ctx.state).ok());

auto shared_state = init_sink_and_source(sink_op, source_op, ctx);
init_topn_runtime_predicate(ctx.state, sink_op->node_id(), false);

Block first_block {ColumnHelper::create_column_with_name<DataTypeInt64>({1, 2, 3, 4}),
ColumnHelper::create_column_with_name<DataTypeInt64>({1, 2, 3, 4})};
ASSERT_TRUE(sink_op->sink(&ctx.state, &first_block, false).ok());
EXPECT_EQ(get_topn_runtime_predicate_value(ctx.state, sink_op->node_id()),
Field::create_field<TYPE_BIGINT>(3));

Block second_block {ColumnHelper::create_column_with_name<DataTypeInt64>({2, 5}),
ColumnHelper::create_column_with_name<DataTypeInt64>({2, 5})};
ASSERT_TRUE(sink_op->sink(&ctx.state, &second_block, true).ok());
EXPECT_EQ(get_topn_runtime_predicate_value(ctx.state, sink_op->node_id()),
Field::create_field<TYPE_BIGINT>(4));
}

TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy,
test_null_boundary_does_not_update_runtime_predicate) {
OperatorContext ctx;
auto sink_op = std::make_shared<MockAggsinkOperator>();
sink_op->_aggregate_evaluators.push_back(create_mock_agg_fn_evaluator(
ctx.pool, MockSlotRef::create_mock_contexts(1, std::make_shared<DataTypeInt64>()),
false, false));
sink_op->_pool = &ctx.pool;
sink_op->_limit = 1;
sink_op->_probe_expr_ctxs = MockSlotRef::create_mock_contexts(
0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>()));
sink_op->_order_directions = {1};
sink_op->_null_directions = {-1};
sink_op->_do_sort_limit = true;
ASSERT_TRUE(sink_op->prepare(&ctx.state).ok());

auto source_op = std::make_shared<MockAggSourceOperator>();
source_op->mock_row_descriptor.reset(new MockRowDescriptor {
{std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>()),
std::make_shared<DataTypeInt64>()},
&ctx.pool});
source_op->_without_key = false;
source_op->_needs_finalize = true;
ASSERT_TRUE(source_op->prepare(&ctx.state).ok());

auto shared_state = init_sink_and_source(sink_op, source_op, ctx);
init_topn_runtime_predicate(ctx.state, sink_op->node_id(), true);
shared_state->reach_limit = true;
shared_state->do_sort_limit = true;
shared_state->limit_columns.emplace_back(
std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())->create_column());
shared_state->limit_columns[0]->insert_default();
shared_state->limit_columns_min = 0;

auto* local_state = static_cast<AggSinkLocalState*>(ctx.state.get_sink_local_state());
ASSERT_TRUE(local_state->_update_runtime_predicate(&ctx.state).ok());
auto& predicate = ctx.state.get_query_ctx()->get_runtime_predicate(sink_op->node_id());
EXPECT_FALSE(predicate.has_value());

shared_state->limit_columns[0]->insert(Field::create_field<TYPE_BIGINT>(7));
shared_state->limit_columns_min = 1;
ASSERT_TRUE(local_state->_update_runtime_predicate(&ctx.state).ok());
ASSERT_TRUE(predicate.has_value());
EXPECT_EQ(predicate.get_value(), Field::create_field<TYPE_BIGINT>(7));
}

TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_need_finalize_with_order_by) {
/*
select column1, sum(column2) from test_table group by column1 order by column1 limit 3;
Expand Down Expand Up @@ -258,6 +369,7 @@ TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_need_finalize_with_order_
EXPECT_TRUE(source_op->prepare(&ctx.state).ok());

auto shared_state = init_sink_and_source(sink_op, source_op, ctx);
init_topn_runtime_predicate(ctx.state, sink_op->node_id(), true);

{
Block block {ColumnHelper::create_column_with_name<DataTypeInt64>({1, 2, 3, 4, 5, 6}),
Expand All @@ -266,6 +378,8 @@ TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_need_finalize_with_order_
std::cout << block.dump_data() << std::endl;
auto st = sink_op->sink(&ctx.state, &block, false);
EXPECT_TRUE(st.ok()) << st.msg();
EXPECT_EQ(get_topn_runtime_predicate_value(ctx.state, sink_op->node_id()),
Field::create_field<TYPE_BIGINT>(3));
}

{
Expand All @@ -275,6 +389,8 @@ TEST_F(AggOperatorGroupByLimitOptTestWithGroupBy, test_need_finalize_with_order_
std::cout << block.dump_data() << std::endl;
auto st = sink_op->sink(&ctx.state, &block, true);
EXPECT_TRUE(st.ok()) << st.msg();
EXPECT_EQ(get_topn_runtime_predicate_value(ctx.state, sink_op->node_id()),
Field::create_field<TYPE_BIGINT>(2));
}

{
Expand Down
Loading
Loading