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 <[email protected]> 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 <[email protected]> --- 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: [email protected] For additional commands, e-mail: [email protected]
