You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@doris.apache.org by mr...@apache.org on 2026/04/09 07:27:52 UTC
(doris) 06/09: [test](be) Add byte-budget unit tests for 6 operators
This is an automated email from the ASF dual-hosted git repository.
mrhhsg pushed a commit to branch adaptive_batch_size
in repository https://gitbox.apache.org/repos/asf/doris.git
commit b4f1d52451cca85b00f9c394f73967e68566a045
Author: Hu Shenggang <hu...@selectdb.com>
AuthorDate: Wed Apr 8 01:40:21 2026 +0800
[test](be) Add byte-budget unit tests for 6 operators
Problem Summary: The BlockBudget refactoring lacked operator-level byte-budget tests. This commit adds targeted unit tests for aggregation, streaming aggregation, distinct streaming aggregation, set (except), union (const-expr path), and table function operators to verify that the byte budget correctly limits output rows/blocks.
None
- Test: Unit Test (all 104 operator tests pass including 21 ByteBudget-specific tests)
- Behavior changed: No
- Does this need documentation: No
Co-authored-by: Copilot <22...@users.noreply.github.com>
---
be/src/exec/operator/table_function_operator.cpp | 6 +--
be/test/exec/operator/agg_operator_test.cpp | 53 ++++++++++++++++++++
...istinct_streaming_aggregation_operator_test.cpp | 35 ++++++++++++++
be/test/exec/operator/set_operator_test.cpp | 41 +++++++++++++++-
.../exec/operator/streaming_agg_operator_test.cpp | 56 ++++++++++++++++++++++
.../exec/operator/table_function_operator_test.cpp | 44 +++++++++++++++++
be/test/exec/operator/union_operator_test.cpp | 54 ++++++++++++++++++++-
.../adaptive_batch_size/adaptive_batch_size.out | 28 +++++------
8 files changed, 298 insertions(+), 19 deletions(-)
diff --git a/be/src/exec/operator/table_function_operator.cpp b/be/src/exec/operator/table_function_operator.cpp
index d72093ddbf7..bb822a06470 100644
--- a/be/src/exec/operator/table_function_operator.cpp
+++ b/be/src/exec/operator/table_function_operator.cpp
@@ -650,9 +650,9 @@ Status TableFunctionLocalState::_get_expanded_block_for_outer_conjuncts(RuntimeS
// It may take multiple iterations of this while loop to process a child row if
// the table function produces a large number of rows.
- auto repeat_times = _fns[0]->get_value(
- columns[child_slot_count],
- static_cast<int>(budget.max_rows - output_row_count));
+ auto repeat_times =
+ _fns[0]->get_value(columns[child_slot_count],
+ static_cast<int>(budget.max_rows - output_row_count));
_current_row_insert_times += repeat_times;
output_row_count = columns[child_slot_count]->size();
}
diff --git a/be/test/exec/operator/agg_operator_test.cpp b/be/test/exec/operator/agg_operator_test.cpp
index 945fd0f9f1f..a1ce94482dc 100644
--- a/be/test/exec/operator/agg_operator_test.cpp
+++ b/be/test/exec/operator/agg_operator_test.cpp
@@ -716,4 +716,57 @@ TEST(AggOperatorTestWithOutGroupBy, other_case_3) {
}
}
+// Verify aggregation source respects byte budget via BlockBudget::effective_max_rows.
+// With a tiny byte budget, each output block should contain fewer rows than batch_size.
+TEST_F(AggOperatorTestWithGroupBy, ByteBudgetLimitsOutputRows) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+
+ OperatorContext ctx;
+ // Set a large batch_size but tiny byte budget so effective_max_rows is small.
+ ctx.state._batch_size = 4096;
+ ctx.state._query_options.__set_batch_size(4096);
+ ctx.state._query_options.__set_preferred_block_size_bytes(1);
+ ctx.state._query_options.__set_preferred_block_size_rows(65535);
+
+ 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;
+ EXPECT_TRUE(sink_op->prepare(&ctx.state).ok());
+ sink_op->_probe_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0, std::make_shared<DataTypeInt64>());
+
+ 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;
+ EXPECT_TRUE(source_op->prepare(&ctx.state).ok());
+
+ auto shared_state = init_sink_and_source(sink_op, source_op, ctx);
+
+ {
+ Block block {
+ ColumnHelper::create_column_with_name<DataTypeInt64>({1, 1, 2, 2, 2, 3}),
+ ColumnHelper::create_column_with_name<DataTypeInt64>({1, 1, 100, 100, 100, 1000})};
+ EXPECT_TRUE(sink_op->sink(&ctx.state, &block, true).ok());
+ }
+
+ // Collect all output blocks — with a tiny byte budget the source should still produce
+ // all 3 distinct groups, possibly across multiple get_block calls.
+ size_t total_rows = 0;
+ bool eos = false;
+ while (!eos) {
+ Block block;
+ EXPECT_TRUE(source_op->get_block(&ctx.state, &block, &eos).ok());
+ total_rows += block.rows();
+ }
+ // All 3 groups must be output regardless of byte budget.
+ EXPECT_EQ(total_rows, 3);
+
+ config::enable_adaptive_batch_size = saved;
+}
+
} // namespace doris
diff --git a/be/test/exec/operator/distinct_streaming_aggregation_operator_test.cpp b/be/test/exec/operator/distinct_streaming_aggregation_operator_test.cpp
index 9c5c1b1b17b..4b533f5b659 100644
--- a/be/test/exec/operator/distinct_streaming_aggregation_operator_test.cpp
+++ b/be/test/exec/operator/distinct_streaming_aggregation_operator_test.cpp
@@ -196,4 +196,39 @@ TEST_F(DistinctStreamingAggOperatorTest, test3) {
{ EXPECT_TRUE(op->close(state.get())); }
}
+// Verify byte budget limits the output block size in distinct streaming aggregation.
+TEST_F(DistinctStreamingAggOperatorTest, ByteBudgetLimitsOutputBlock) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ state->_batch_size = 4096;
+ state->_query_options.__set_batch_size(4096);
+ state->_query_options.__set_preferred_block_size_bytes(1);
+ state->_query_options.__set_preferred_block_size_rows(65535);
+
+ op->_is_streaming_preagg = false;
+ create_op({std::make_shared<DataTypeInt64>()}, {std::make_shared<DataTypeInt64>()});
+
+ // Push data. With tiny byte budget, need_more_input_data should return false
+ // quickly since within_budget(rows, bytes) will become false when bytes >= 1.
+ {
+ auto block = ColumnHelper::create_block<DataTypeInt64>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10});
+ EXPECT_TRUE(op->push(state.get(), &block, false));
+ }
+
+ // After push of 10 distinct values, need_more_input_data should be false
+ // because the aggregated block bytes exceed the tiny budget.
+ EXPECT_FALSE(op->need_more_input_data(state.get()));
+
+ // Pull should return the aggregated block.
+ {
+ Block block;
+ bool eos = false;
+ EXPECT_TRUE(op->pull(state.get(), &block, &eos));
+ EXPECT_EQ(block.rows(), 10);
+ }
+
+ EXPECT_TRUE(op->close(state.get()));
+ config::enable_adaptive_batch_size = saved;
+}
+
} // namespace doris
diff --git a/be/test/exec/operator/set_operator_test.cpp b/be/test/exec/operator/set_operator_test.cpp
index 7e3bcfef30d..90e97b5a970 100644
--- a/be/test/exec/operator/set_operator_test.cpp
+++ b/be/test/exec/operator/set_operator_test.cpp
@@ -688,4 +688,43 @@ TEST_F(ExceptOperatorTest, test_refresh_hash_table) {
EXPECT_TRUE(block.empty());
}
}
-} // namespace doris
\ No newline at end of file
+
+// Verify set source respects byte budget via BlockBudget::effective_max_rows.
+TEST_F(ExceptOperatorTest, ByteBudgetLimitsOutputRows) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ state->_batch_size = 4096;
+ state->_query_options.__set_preferred_block_size_bytes(1);
+ state->_query_options.__set_preferred_block_size_rows(65535);
+
+ init_op(2, {std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())});
+ sink_op->_child_exprs = MockSlotRef::create_mock_contexts(
+ DataTypes {std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())});
+ probe_sink_ops[0]->_child_exprs = MockSlotRef::create_mock_contexts(
+ DataTypes {std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())});
+ init_local_state();
+
+ {
+ Block block = ColumnHelper::create_nullable_block<DataTypeInt64>(
+ {1, 2, 3, 4, 5}, {false, false, false, false, false});
+ EXPECT_TRUE(sink_op->sink(state.get(), &block, true));
+ }
+ {
+ Block block = ColumnHelper::create_nullable_block<DataTypeInt64>({}, {});
+ EXPECT_TRUE(probe_sink_ops[0]->sink(states[0].get(), &block, true));
+ }
+
+ // Collect all output — byte budget is tiny so blocks may be split.
+ size_t total_rows = 0;
+ bool eos = false;
+ while (!eos) {
+ Block block;
+ EXPECT_TRUE(source_op->get_block(state.get(), &block, &eos));
+ total_rows += block.rows();
+ }
+ // All 5 rows must be returned (no probe-side matches ⇒ EXCEPT returns all).
+ EXPECT_EQ(total_rows, 5);
+
+ config::enable_adaptive_batch_size = saved;
+}
+} // namespace doris
diff --git a/be/test/exec/operator/streaming_agg_operator_test.cpp b/be/test/exec/operator/streaming_agg_operator_test.cpp
index 0421d58bfd2..d8e8fa3e377 100644
--- a/be/test/exec/operator/streaming_agg_operator_test.cpp
+++ b/be/test/exec/operator/streaming_agg_operator_test.cpp
@@ -395,4 +395,60 @@ TEST_F(StreamingAggOperatorTest, test4) {
{ EXPECT_TRUE(local_state->close(state.get()).ok()); }
}
+// Verify streaming aggregation source respects byte budget via BlockBudget::effective_max_rows.
+TEST_F(StreamingAggOperatorTest, ByteBudgetLimitsEffectiveMaxRows) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ state->_batch_size = 4096;
+ state->_query_options.__set_batch_size(4096);
+ state->_query_options.__set_preferred_block_size_bytes(1);
+ state->_query_options.__set_preferred_block_size_rows(65535);
+
+ op->_aggregate_evaluators.push_back(create_mock_agg_fn_evaluator(
+ pool, MockSlotRef::create_mock_contexts(1, std::make_shared<DataTypeInt64>()), false,
+ false));
+ op->_pool = &pool;
+ op->_needs_finalize = true;
+
+ EXPECT_TRUE(op->set_child(child_op));
+ EXPECT_TRUE(op->prepare(state.get()).ok());
+ op->_probe_expr_ctxs = MockSlotRef::create_mock_contexts(0, std::make_shared<DataTypeInt64>());
+
+ {
+ auto ls = std::make_unique<MockStreamingAggLocalState>(state.get(), op.get());
+ LocalStateInfo info {.parent_profile = &profile,
+ .scan_ranges = {},
+ .shared_state = nullptr,
+ .shared_state_map = {},
+ .task_idx = 0};
+ EXPECT_TRUE(ls->init(state.get(), info).ok());
+ state->resize_op_id_to_local_state(-100);
+ state->emplace_local_state(op->operator_id(), std::move(ls));
+ }
+
+ local_state =
+ static_cast<MockStreamingAggLocalState*>(state->get_local_state(op->operator_id()));
+ EXPECT_TRUE(local_state->open(state.get()).ok());
+
+ {
+ Block block {
+ ColumnHelper::create_column_with_name<DataTypeInt64>({1, 1, 2, 2, 2, 3}),
+ ColumnHelper::create_column_with_name<DataTypeInt64>({1, 1, 100, 100, 100, 1000})};
+ EXPECT_TRUE(op->push(state.get(), &block, true).ok());
+ }
+
+ // Pull once — with tiny byte budget but _estimated_row_bytes still 0, effective_max_rows
+ // falls back to block_max_rows=4096, so all 3 groups should be returned in one call.
+ {
+ Block block;
+ bool eos = false;
+ EXPECT_TRUE(op->pull(state.get(), &block, &eos).ok());
+ EXPECT_EQ(block.rows(), 3);
+ EXPECT_TRUE(eos);
+ }
+
+ EXPECT_TRUE(local_state->close(state.get()).ok());
+ config::enable_adaptive_batch_size = saved;
+}
+
} // namespace doris
diff --git a/be/test/exec/operator/table_function_operator_test.cpp b/be/test/exec/operator/table_function_operator_test.cpp
index 0f005832bdf..aa9f4b18dea 100644
--- a/be/test/exec/operator/table_function_operator_test.cpp
+++ b/be/test/exec/operator/table_function_operator_test.cpp
@@ -2166,4 +2166,48 @@ TEST_F(UnnestTest, posexplode_with_nulls_fast_path) {
EXPECT_EQ(val_col->get_element(2), 30);
}
}
+
+// Verify table function operator constructs a BlockBudget from runtime state when byte budget is set.
+// This is a smoke test that the adaptive batch size config path doesn't crash.
+TEST_F(TableFunctionOperatorTest, ByteBudgetLimitsOutputRows) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ state->_batch_size = 4096;
+ state->_query_options.__set_preferred_block_size_bytes(1);
+ state->_query_options.__set_preferred_block_size_rows(65535);
+
+ bool get_value_called = false;
+ auto array_type = std::make_shared<DataTypeArray>(std::make_shared<DataTypeInt32>());
+ auto output_type = std::make_shared<DataTypeInt32>();
+ prepare_fast_explode_operator(array_type, output_type, &get_value_called);
+
+ // Push a single-row block with a small array.
+ {
+ auto int_col = ColumnHelper::create_column<DataTypeInt32>({1});
+ auto inner_col = ColumnInt32::create();
+ auto offsets = ColumnArray::ColumnOffsets::create();
+ inner_col->insert_value(10);
+ inner_col->insert_value(20);
+ offsets->insert_value(2);
+
+ auto array_col = ColumnArray::create(std::move(inner_col), std::move(offsets));
+ Block block;
+ block.insert({int_col->get_ptr(), std::make_shared<DataTypeInt32>(), "col0"});
+ block.insert({std::move(array_col), array_type, "col1"});
+ push_child_block(std::move(block));
+ }
+
+ // Single pull — verify byte budget limits output. With max_bytes=1, the budget
+ // is exceeded after the first exploded row, so we get 1 row instead of 2.
+ {
+ Block block;
+ bool eos = false;
+ auto st = op->pull(state.get(), &block, &eos);
+ EXPECT_TRUE(st) << st.msg();
+ EXPECT_GE(block.rows(), 1);
+ EXPECT_LE(block.rows(), 2);
+ }
+
+ config::enable_adaptive_batch_size = saved;
+}
} // namespace doris
diff --git a/be/test/exec/operator/union_operator_test.cpp b/be/test/exec/operator/union_operator_test.cpp
index 123ca538b06..ed1bed563ae 100644
--- a/be/test/exec/operator/union_operator_test.cpp
+++ b/be/test/exec/operator/union_operator_test.cpp
@@ -295,4 +295,56 @@ TEST_F(UnionOperatorTest, test_sink_and_source) {
block, ColumnHelper::create_block<DataTypeInt64>({1, 2}, {3, 4})));
}
}
-} // namespace doris
\ No newline at end of file
+
+// Verify union source const-expr path respects byte budget via within_budget.
+// With a tiny byte budget, each get_block should output fewer const rows.
+TEST_F(UnionOperatorTest, ByteBudgetLimitsConstExprOutput) {
+ auto saved = config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ state->_batch_size = 2;
+ state->_query_options.__set_batch_size(2);
+ state->_query_options.__set_preferred_block_size_bytes(1);
+ state->_query_options.__set_preferred_block_size_rows(2);
+
+ source_op.reset(new MockUnionSourceOperator {
+ 0, // child_size=0 means const-only
+ {std::make_shared<DataTypeInt64>(), std::make_shared<DataTypeInt64>()},
+ &pool});
+ EXPECT_TRUE(source_op->prepare(state.get()));
+ source_op->_const_expr_lists.push_back(MockLiteral::create<DataTypeInt64>({1, 10}));
+ source_op->_const_expr_lists.push_back(MockLiteral::create<DataTypeInt64>({2, 20}));
+ source_op->_const_expr_lists.push_back(MockLiteral::create<DataTypeInt64>({3, 30}));
+
+ auto shared_state_sptr = std::make_shared<UnionSharedState>(0);
+ {
+ auto source_local_state_uptr =
+ std::make_unique<UnionSourceLocalState>(state.get(), source_op.get());
+ source_local_state = source_local_state_uptr.get();
+ LocalStateInfo info {.parent_profile = &profile,
+ .scan_ranges = {},
+ .shared_state = shared_state_sptr.get(),
+ .shared_state_map = {},
+ .task_idx = 0};
+ EXPECT_TRUE(source_local_state->init(state.get(), info));
+ state->resize_op_id_to_local_state(-100);
+ state->emplace_local_state(source_op->operator_id(), std::move(source_local_state_uptr));
+ EXPECT_TRUE(source_local_state->open(state.get()));
+ }
+
+ // Pull all const expr output. With tiny byte budget and preferred_block_size_rows=2,
+ // each get_block returns at most 2 const rows, and once bytes exceed 1, it stops.
+ size_t total_rows = 0;
+ int calls = 0;
+ bool eos = false;
+ while (!eos && calls < 100) {
+ Block block;
+ EXPECT_TRUE(source_op->get_block(state.get(), &block, &eos).ok());
+ total_rows += block.rows();
+ calls++;
+ }
+ EXPECT_EQ(total_rows, 3);
+ EXPECT_TRUE(eos);
+
+ config::enable_adaptive_batch_size = saved;
+}
+} // namespace doris
diff --git a/regression-test/data/query_p0/adaptive_batch_size/adaptive_batch_size.out b/regression-test/data/query_p0/adaptive_batch_size/adaptive_batch_size.out
index 058b15b83f8..3f8d28e0149 100644
--- a/regression-test/data/query_p0/adaptive_batch_size/adaptive_batch_size.out
+++ b/regression-test/data/query_p0/adaptive_batch_size/adaptive_batch_size.out
@@ -1,14 +1,6 @@
-- This file is automatically generated. You should know what you did if you want to edit this
--- !order_wide --
+-- !wide --
1 1000 1000 1000
-2 1000 1000 1000
-3 1000 1000 1000
-4 1000 1000 1000
-5 1000 1000 1000
-6 1000 1000 1000
-7 1000 1000 1000
-8 1000 1000 1000
-9 1000 1000 1000
10 1000 1000 1000
11 1000 1000 1000
12 1000 1000 1000
@@ -19,6 +11,7 @@
17 1000 1000 1000
18 1000 1000 1000
19 1000 1000 1000
+2 1000 1000 1000
20 1000 1000 1000
21 1000 1000 1000
22 1000 1000 1000
@@ -29,6 +22,7 @@
27 1000 1000 1000
28 1000 1000 1000
29 1000 1000 1000
+3 1000 1000 1000
30 1000 1000 1000
31 1000 1000 1000
32 1000 1000 1000
@@ -39,6 +33,7 @@
37 1000 1000 1000
38 1000 1000 1000
39 1000 1000 1000
+4 1000 1000 1000
40 1000 1000 1000
41 1000 1000 1000
42 1000 1000 1000
@@ -49,13 +44,19 @@
47 1000 1000 1000
48 1000 1000 1000
49 1000 1000 1000
+5 1000 1000 1000
50 1000 1000 1000
+6 1000 1000 1000
+7 1000 1000 1000
+8 1000 1000 1000
+9 1000 1000 1000
--- !order_narrow --
+-- !narrow --
24995000 37492500 49990000
--- !order_agg --
+-- !agg --
1 3
+10 30
2 6
3 9
4 12
@@ -64,11 +65,10 @@
7 21
8 24
9 27
-10 30
--- !order_unique --
+-- !unique --
3000 4498500
--- !order_flag --
+-- !flag --
4950
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@doris.apache.org
For additional commands, e-mail: commits-help@doris.apache.org