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 741616fb778457a688e6f3dde752d5fc6d7888cc Author: Hu Shenggang <[email protected]> AuthorDate: Tue Apr 7 22:51:41 2026 +0800 [refactor](be) Extract BlockBudget utility for block row/byte budget control ### What problem does this PR solve? Issue Number: close #xxx Related PR: #61535 Problem Summary: Multiple operators (aggregation_source, streaming_aggregation, set_source, distinct_streaming_aggregation, union_source, table_function, nested_loop_join_probe, schema_scan) had identical copy-pasted logic for computing effective max rows from byte budgets, checking within-budget conditions, and computing remaining row capacity. This duplication increases maintenance burden and inconsistency risk. Introduce `BlockBudget` struct in `be/src/common/block_budget.h` that centralizes the 4 common patterns: - `effective_max_rows(estimated_row_bytes)` — Pattern A (5 call sites) - `within_budget(rows, bytes)` — Pattern B (6 call sites) - `exceeded(rows, bytes)` — break conditions - `remaining_rows(rows, bytes)` — Pattern C (2 call sites) Net effect: -33 lines across 8 operator files, replaced with concise single-line calls. All 59 related unit tests pass. ### Release note None ### Check List (For Author) - Test: Unit Test (59 tests including 14 new BlockBudget tests) - Behavior changed: No - Does this need documentation: No Co-authored-by: Copilot <[email protected]> --- be/src/common/block_budget.h | 93 +++++++++++++++ .../exec/operator/aggregation_source_operator.cpp | 17 +-- .../distinct_streaming_aggregation_operator.cpp | 31 ++--- .../operator/nested_loop_join_probe_operator.cpp | 31 ++--- be/src/exec/operator/schema_scan_operator.cpp | 6 +- be/src/exec/operator/set_source_operator.cpp | 9 +- .../operator/streaming_aggregation_operator.cpp | 9 +- be/src/exec/operator/table_function_operator.cpp | 20 ++-- be/src/exec/operator/union_source_operator.cpp | 6 +- be/test/common/block_budget_test.cpp | 132 +++++++++++++++++++++ 10 files changed, 273 insertions(+), 81 deletions(-) diff --git a/be/src/common/block_budget.h b/be/src/common/block_budget.h new file mode 100644 index 00000000000..509e751043b --- /dev/null +++ b/be/src/common/block_budget.h @@ -0,0 +1,93 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include <algorithm> +#include <cstddef> + +namespace doris { + +// Lightweight value type that captures the dual row+byte budget for block +// output sizing. Every operator that needs to respect the adaptive batch +// size feature can construct a BlockBudget from RuntimeState and use the +// helper methods instead of reimplementing the same row/byte logic inline. +// +// Typical usage: +// BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); +// size_t eff = budget.effective_max_rows(estimated_row_bytes); +// while (budget.within_budget(block.rows(), block.bytes())) { ... } +// +struct BlockBudget { + size_t max_rows; + size_t max_bytes; // 0 = byte budget disabled + + BlockBudget(size_t max_rows_, size_t max_bytes_) : max_rows(max_rows_), max_bytes(max_bytes_) {} + + bool has_byte_budget() const { return max_bytes > 0; } + + // Pre-compute effective row limit from an estimated average row byte size. + // When byte budget is disabled or estimated_row_bytes == 0, returns max_rows. + // Always returns at least 1. + size_t effective_max_rows(size_t estimated_row_bytes) const { + if (max_bytes > 0 && estimated_row_bytes > 0) { + size_t bytes_limit = max_bytes / estimated_row_bytes; + return std::max(size_t(1), std::min(max_rows, bytes_limit)); + } + return max_rows; + } + + // Check if a block with the given rows/bytes is still within budget. + // Use this in loop *continuation* conditions (while/for). + bool within_budget(size_t rows, size_t bytes) const { + return rows < max_rows && (max_bytes == 0 || bytes < max_bytes); + } + + // Check if a block with the given rows/bytes has exceeded the budget. + // Use this in loop *break* conditions. + bool exceeded(size_t rows, size_t bytes) const { + return rows >= max_rows || (max_bytes > 0 && bytes >= max_bytes); + } + + // Compute how many more rows can be added to a block that currently + // has current_rows rows and current_bytes bytes, respecting both the + // row cap and the byte budget. + // The 3-arg overload accepts an explicit estimated_row_bytes (useful when + // the estimate comes from a different source, e.g. a child block). + // The 2-arg overload derives the estimate from current_bytes / current_rows. + // Returns 0 when the block is already at or over budget. + size_t remaining_rows(size_t current_rows, size_t current_bytes, + size_t estimated_row_bytes) const { + size_t row_capacity = (current_rows < max_rows) ? (max_rows - current_rows) : 0; + if (max_bytes > 0 && estimated_row_bytes > 0) { + if (current_bytes >= max_bytes) { + return 0; + } + size_t byte_capacity = (max_bytes - current_bytes) / estimated_row_bytes; + row_capacity = std::min(row_capacity, byte_capacity); + } + return row_capacity; + } + + size_t remaining_rows(size_t current_rows, size_t current_bytes) const { + size_t estimated = + (current_rows > 0 && current_bytes > 0) ? (current_bytes / current_rows) : 0; + return remaining_rows(current_rows, current_bytes, estimated); + } +}; + +} // namespace doris diff --git a/be/src/exec/operator/aggregation_source_operator.cpp b/be/src/exec/operator/aggregation_source_operator.cpp index 1fe67c934a7..0535daf18bc 100644 --- a/be/src/exec/operator/aggregation_source_operator.cpp +++ b/be/src/exec/operator/aggregation_source_operator.cpp @@ -20,6 +20,7 @@ #include <memory> #include <string> +#include "common/block_budget.h" #include "common/exception.h" #include "core/column/column_fixed_length_object.h" #include "exec/operator/operator.h" @@ -122,12 +123,8 @@ Status AggLocalState::_get_results_with_serialized_key(RuntimeState* state, Bloc } // Compute effective max rows based on estimated bytes per row. - size_t effective_max_rows = state->block_max_rows(); - const size_t block_max_bytes = state->block_max_bytes(); - if (block_max_bytes > 0 && _estimated_row_bytes > 0) { - size_t bytes_limit = block_max_bytes / _estimated_row_bytes; - effective_max_rows = std::max(size_t(1), std::min(effective_max_rows, bytes_limit)); - } + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + const size_t effective_max_rows = budget.effective_max_rows(_estimated_row_bytes); std::visit( Overload { @@ -322,12 +319,8 @@ Status AggLocalState::_get_with_serialized_key_result(RuntimeState* state, Block SCOPED_TIMER(_get_results_timer); // Compute effective max rows based on estimated bytes per row. - size_t effective_max_rows = state->block_max_rows(); - const size_t block_max_bytes = state->block_max_bytes(); - if (block_max_bytes > 0 && _estimated_row_bytes > 0) { - size_t bytes_limit = block_max_bytes / _estimated_row_bytes; - effective_max_rows = std::max(size_t(1), std::min(effective_max_rows, bytes_limit)); - } + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + const size_t effective_max_rows = budget.effective_max_rows(_estimated_row_bytes); std::visit( Overload { diff --git a/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp b/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp index 64aa5f1acc0..5a0d25b4924 100644 --- a/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp +++ b/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp @@ -22,6 +22,7 @@ #include <memory> #include <utility> +#include "common/block_budget.h" #include "common/compiler_util.h" // IWYU pragma: keep #include "exec/operator/streaming_agg_min_reduction.h" #include "exprs/vectorized_agg_fn.h" @@ -210,21 +211,9 @@ Status DistinctStreamingAggLocalState::_distinct_pre_agg_with_serialized_key( DCHECK_EQ(_cache_block.rows(), 0); // Compute maximum additional rows for out_block respecting both // row limit (batch_size) and byte budget (block_max_bytes). - size_t max_rows_to_add = (out_block->rows() < (size_t)batch_size) - ? ((size_t)batch_size - out_block->rows()) - : 0; - if (block_max_bytes > 0 && out_block->rows() > 0) { - size_t current_bytes = out_block->bytes(); - if (current_bytes >= block_max_bytes) { - max_rows_to_add = 0; - } else { - size_t avg_row_bytes = current_bytes / out_block->rows(); - if (avg_row_bytes > 0) { - size_t rows_for_budget = (block_max_bytes - current_bytes) / avg_row_bytes; - max_rows_to_add = std::min(max_rows_to_add, rows_for_budget); - } - } - } + const BlockBudget budget(batch_size, block_max_bytes); + const size_t max_rows_to_add = + budget.remaining_rows(out_block->rows(), out_block->bytes()); if (_distinct_row.size() > max_rows_to_add) { for (int i = 0; i < key_size; ++i) { auto output_dst = out_block->get_by_position(i).column->assume_mutable(); @@ -408,13 +397,11 @@ Status DistinctStreamingAggOperatorX::pull(RuntimeState* state, Block* block, bo bool DistinctStreamingAggOperatorX::need_more_input_data(RuntimeState* state) const { auto& local_state = get_local_state(state); - const auto block_max_bytes = state->block_max_bytes(); - const bool need_batch = - local_state._stop_emplace_flag - ? local_state._aggregated_block->empty() - : (local_state._aggregated_block->rows() < state->block_max_rows() && - (block_max_bytes == 0 || - local_state._aggregated_block->bytes() < block_max_bytes)); + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + const bool need_batch = local_state._stop_emplace_flag + ? local_state._aggregated_block->empty() + : budget.within_budget(local_state._aggregated_block->rows(), + local_state._aggregated_block->bytes()); return need_batch && !(local_state._child_eos || local_state._reach_limit); } diff --git a/be/src/exec/operator/nested_loop_join_probe_operator.cpp b/be/src/exec/operator/nested_loop_join_probe_operator.cpp index a0cc4bf25ed..73fa295d048 100644 --- a/be/src/exec/operator/nested_loop_join_probe_operator.cpp +++ b/be/src/exec/operator/nested_loop_join_probe_operator.cpp @@ -19,6 +19,7 @@ #include <memory> +#include "common/block_budget.h" #include "common/cast_set.h" #include "common/exception.h" #include "core/block/block.h" @@ -195,7 +196,8 @@ void NestedLoopJoinProbeLocalState::_generate_block_base_probe(RuntimeState* sta }; const auto block_max_bytes = state->block_max_bytes(); - size_t effective_max_rows = state->block_max_rows(); + const BlockBudget budget(state->block_max_rows(), block_max_bytes); + size_t effective_max_rows = budget.max_rows; bool bytes_estimated = false; while (_join_block.rows() + add_rows() <= effective_max_rows) { while (_current_build_pos == _shared_state->build_blocks.size() || @@ -245,9 +247,8 @@ void NestedLoopJoinProbeLocalState::_generate_block_base_probe(RuntimeState* sta } } - DCHECK_LE(_join_block.rows(), state->block_max_rows()) - << "join block rows:" << _join_block.rows() - << ", state batch size:" << state->block_max_rows() + DCHECK_LE(_join_block.rows(), budget.max_rows) + << "join block rows:" << _join_block.rows() << ", state batch size:" << budget.max_rows << "probe_block rows:" << probe_block->rows() << " build blocks size:" << _shared_state->build_blocks.size(); } @@ -283,7 +284,8 @@ void NestedLoopJoinProbeLocalState::_generate_block_base_build(RuntimeState* sta } const auto block_max_bytes = state->block_max_bytes(); - size_t effective_max_rows = state->block_max_rows(); + const BlockBudget budget(state->block_max_rows(), block_max_bytes); + size_t effective_max_rows = budget.max_rows; bool bytes_estimated = false; while (_join_block.rows() + probe_rows <= effective_max_rows) { // The current build row has processed the entire probe block; move to the next build row @@ -321,17 +323,13 @@ void NestedLoopJoinProbeLocalState::_generate_block_base_build(RuntimeState* sta // Refine row limit based on actual bytes per row after first data-producing iteration if (!bytes_estimated && block_max_bytes > 0 && _join_block.rows() > 0) { bytes_estimated = true; - size_t bytes_per_row = _join_block.bytes() / _join_block.rows(); - if (bytes_per_row > 0) { - effective_max_rows = std::min(effective_max_rows, - std::max(size_t(1), block_max_bytes / bytes_per_row)); - } + effective_max_rows = + budget.effective_max_rows(_join_block.bytes() / _join_block.rows()); } } - DCHECK_LE(_join_block.rows(), state->block_max_rows()) - << "join block rows:" << _join_block.rows() - << ", state batch size:" << state->block_max_rows() + DCHECK_LE(_join_block.rows(), budget.max_rows) + << "join block rows:" << _join_block.rows() << ", state batch size:" << budget.max_rows << "probe_block rows:" << probe_block->rows() << "build block rows:" << build_block.rows(); } @@ -437,11 +435,8 @@ void NestedLoopJoinProbeLocalState::_finalize_current_phase(Block& block, size_t for (const auto& col : dst_columns) { current_bytes += col->byte_size(); } - size_t bytes_per_row = current_bytes / column_size; - if (bytes_per_row > 0) { - effective_batch_size = - std::min(batch_size, std::max(size_t(1), block_max_bytes / bytes_per_row)); - } + const BlockBudget budget(batch_size, block_max_bytes); + effective_batch_size = budget.effective_max_rows(current_bytes / column_size); } if constexpr (BuildSide) { DCHECK(!p._is_mark_join); diff --git a/be/src/exec/operator/schema_scan_operator.cpp b/be/src/exec/operator/schema_scan_operator.cpp index 128cadad46e..aec72a3cffe 100644 --- a/be/src/exec/operator/schema_scan_operator.cpp +++ b/be/src/exec/operator/schema_scan_operator.cpp @@ -21,6 +21,7 @@ #include <memory> +#include "common/block_budget.h" #include "core/data_type/data_type_factory.hpp" #include "exec/operator/operator.h" #include "runtime/runtime_profile.h" @@ -243,9 +244,8 @@ Status SchemaScanOperatorX::get_block(RuntimeState* state, Block* block, bool* e break; } - const auto block_max_bytes = state->block_max_bytes(); - if (src_block.rows() >= state->block_max_rows() || - (block_max_bytes > 0 && src_block.bytes() >= block_max_bytes)) { + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + if (budget.exceeded(src_block.rows(), src_block.bytes())) { break; } } diff --git a/be/src/exec/operator/set_source_operator.cpp b/be/src/exec/operator/set_source_operator.cpp index 383a085fcb2..5a36a73cafb 100644 --- a/be/src/exec/operator/set_source_operator.cpp +++ b/be/src/exec/operator/set_source_operator.cpp @@ -20,6 +20,7 @@ #include <memory> #include <type_traits> +#include "common/block_budget.h" #include "common/status.h" #include "exec/operator/operator.h" #include "runtime/runtime_profile.h" @@ -82,12 +83,8 @@ Status SetSourceOperatorX<is_intersect>::get_block(RuntimeState* state, Block* b SCOPED_PEAK_MEM(&local_state._estimate_memory_usage); // Compute effective batch size based on estimated bytes per row. - size_t effective_batch_size = state->block_max_rows(); - const size_t block_max_bytes = state->block_max_bytes(); - if (block_max_bytes > 0 && local_state._estimated_row_bytes > 0) { - size_t bytes_limit = block_max_bytes / local_state._estimated_row_bytes; - effective_batch_size = std::max(size_t(1), std::min(effective_batch_size, bytes_limit)); - } + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + const size_t effective_batch_size = budget.effective_max_rows(local_state._estimated_row_bytes); _create_mutable_cols(local_state, block); { diff --git a/be/src/exec/operator/streaming_aggregation_operator.cpp b/be/src/exec/operator/streaming_aggregation_operator.cpp index 479347e1e92..329c6331272 100644 --- a/be/src/exec/operator/streaming_aggregation_operator.cpp +++ b/be/src/exec/operator/streaming_aggregation_operator.cpp @@ -22,6 +22,7 @@ #include <memory> #include <utility> +#include "common/block_budget.h" #include "common/cast_set.h" #include "common/compiler_util.h" // IWYU pragma: keep #include "core/column/column_fixed_length_object.h" @@ -470,12 +471,8 @@ Status StreamingAggLocalState::_get_results_with_serialized_key(RuntimeState* st } // Compute effective max rows based on estimated bytes per row from previous calls. - size_t effective_max_rows = state->block_max_rows(); - const size_t block_max_bytes = state->block_max_bytes(); - if (block_max_bytes > 0 && _estimated_row_bytes > 0) { - size_t bytes_limit = block_max_bytes / _estimated_row_bytes; - effective_max_rows = std::max(size_t(1), std::min(effective_max_rows, bytes_limit)); - } + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); + const size_t effective_max_rows = budget.effective_max_rows(_estimated_row_bytes); std::visit( Overload { diff --git a/be/src/exec/operator/table_function_operator.cpp b/be/src/exec/operator/table_function_operator.cpp index 414f2a57bc4..d72093ddbf7 100644 --- a/be/src/exec/operator/table_function_operator.cpp +++ b/be/src/exec/operator/table_function_operator.cpp @@ -23,6 +23,7 @@ #include <memory> #include <numeric> +#include "common/block_budget.h" #include "common/cast_set.h" #include "core/assert_cast.h" #include "core/block/block.h" @@ -530,10 +531,9 @@ Status TableFunctionLocalState::get_expanded_block(RuntimeState* state, Block* o } if (use_slow_path) { - const auto block_max_bytes_slow = state->block_max_bytes(); + const BlockBudget budget_slow(state->block_max_rows(), state->block_max_bytes()); bool skip_child_row = false; - while (columns[p._child_slots.size()]->size() < state->block_max_rows() && - (block_max_bytes_slow == 0 || m_block.bytes() < block_max_bytes_slow)) { + while (budget_slow.within_budget(columns[p._child_slots.size()]->size(), m_block.bytes())) { RETURN_IF_CANCELLED(state); if (_child_block->rows() == 0) { @@ -616,12 +616,10 @@ Status TableFunctionLocalState::_get_expanded_block_for_outer_conjuncts(RuntimeS child_row_to_output_rows_indices.push_back(0); } - const auto block_max_rows = state->block_max_rows(); - const auto block_max_bytes = state->block_max_bytes(); + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); auto output_row_count = columns[child_slot_count]->size(); auto within_block_limit = [&]() -> bool { - return output_row_count < block_max_rows && - (block_max_bytes == 0 || m_block.bytes() < block_max_bytes); + return budget.within_budget(output_row_count, m_block.bytes()); }; while (within_block_limit()) { RETURN_IF_CANCELLED(state); @@ -652,8 +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], - block_max_rows - (int)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(); } @@ -663,8 +662,7 @@ Status TableFunctionLocalState::_get_expanded_block_for_outer_conjuncts(RuntimeS // _cur_child_offset == -1 // 2. output_block reaches batch size // fn maybe or maybe not eos - if (output_row_count >= block_max_rows || - (block_max_bytes > 0 && m_block.bytes() >= block_max_bytes)) { + if (budget.exceeded(output_row_count, m_block.bytes())) { _copy_output_slots(columns, p); handled_row_indices.push_back(_cur_child_offset); child_row_to_output_rows_indices.push_back(output_row_count); diff --git a/be/src/exec/operator/union_source_operator.cpp b/be/src/exec/operator/union_source_operator.cpp index ccf91ca4590..e0390e63576 100644 --- a/be/src/exec/operator/union_source_operator.cpp +++ b/be/src/exec/operator/union_source_operator.cpp @@ -20,6 +20,7 @@ #include <functional> #include <utility> +#include "common/block_budget.h" #include "common/status.h" #include "core/block/block.h" #include "core/block/columns_with_type_and_name.h" @@ -153,10 +154,9 @@ Status UnionSourceOperatorX::get_next_const(RuntimeState* state, Block* block) { MutableBlock mblock = VectorizedUtils::build_mutable_mem_reuse_block(block, row_descriptor()); ColumnsWithTypeAndName tmp_block_columns; - const auto block_max_bytes = state->block_max_bytes(); + const BlockBudget budget(state->block_max_rows(), state->block_max_bytes()); for (; _const_expr_list_idx < _const_expr_lists.size() && - mblock.rows() < state->block_max_rows() && - (block_max_bytes == 0 || mblock.bytes() < block_max_bytes); + budget.within_budget(mblock.rows(), mblock.bytes()); ++_const_expr_list_idx) { int const_expr_lists_size = cast_set<int>(_const_expr_lists[_const_expr_list_idx].size()); if (_const_expr_list_idx && const_expr_lists_size != _const_expr_lists[0].size()) { diff --git a/be/test/common/block_budget_test.cpp b/be/test/common/block_budget_test.cpp new file mode 100644 index 00000000000..8f2dd5a57ed --- /dev/null +++ b/be/test/common/block_budget_test.cpp @@ -0,0 +1,132 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "common/block_budget.h" + +#include <gtest/gtest.h> + +namespace doris { + +class BlockBudgetTest : public ::testing::Test {}; + +// ── effective_max_rows ────────────────────────────────────────────────────── + +TEST_F(BlockBudgetTest, EffectiveMaxRowsNoByteBudget) { + BlockBudget b(4096, 0); + EXPECT_EQ(b.effective_max_rows(100), 4096); + EXPECT_EQ(b.effective_max_rows(0), 4096); +} + +TEST_F(BlockBudgetTest, EffectiveMaxRowsZeroEstimate) { + BlockBudget b(4096, 8 * 1024 * 1024); + // When estimate is 0, fall back to max_rows. + EXPECT_EQ(b.effective_max_rows(0), 4096); +} + +TEST_F(BlockBudgetTest, EffectiveMaxRowsByteLimited) { + // 8 MB budget, 10 KB per row → 819 rows (< 4096 max_rows) + BlockBudget b(4096, 8 * 1024 * 1024); + EXPECT_EQ(b.effective_max_rows(10 * 1024), 819); +} + +TEST_F(BlockBudgetTest, EffectiveMaxRowsRowLimited) { + // 8 MB budget, 10 bytes per row → 838860 rows, but max_rows = 4096 + BlockBudget b(4096, 8 * 1024 * 1024); + EXPECT_EQ(b.effective_max_rows(10), 4096); +} + +TEST_F(BlockBudgetTest, EffectiveMaxRowsReturnsAtLeastOne) { + // Huge rows: 100 MB per row, 8 MB budget → 0, but clamped to 1 + BlockBudget b(4096, 8 * 1024 * 1024); + EXPECT_EQ(b.effective_max_rows(100 * 1024 * 1024), 1); +} + +// ── within_budget / exceeded ──────────────────────────────────────────────── + +TEST_F(BlockBudgetTest, WithinBudgetNoByteBudget) { + BlockBudget b(100, 0); + EXPECT_TRUE(b.within_budget(0, 0)); + EXPECT_TRUE(b.within_budget(99, 999999999)); + EXPECT_FALSE(b.within_budget(100, 0)); + EXPECT_FALSE(b.within_budget(200, 0)); +} + +TEST_F(BlockBudgetTest, WithinBudgetWithByteBudget) { + BlockBudget b(100, 1000); + EXPECT_TRUE(b.within_budget(50, 500)); // both under + EXPECT_FALSE(b.within_budget(100, 500)); // rows hit + EXPECT_FALSE(b.within_budget(50, 1000)); // bytes hit + EXPECT_FALSE(b.within_budget(100, 1000)); // both hit +} + +TEST_F(BlockBudgetTest, ExceededIsInverseOfWithinBudget) { + BlockBudget b(100, 1000); + // Note: exceeded uses >=, within_budget uses <, so they should be + // perfect logical inverses. + for (size_t r : {0, 50, 99, 100, 200}) { + for (size_t bytes : {0, 500, 999, 1000, 2000}) { + EXPECT_EQ(b.exceeded(r, bytes), !b.within_budget(r, bytes)) + << "r=" << r << " bytes=" << bytes; + } + } +} + +// ── remaining_rows ────────────────────────────────────────────────────────── + +TEST_F(BlockBudgetTest, RemainingRowsNoByteBudget) { + BlockBudget b(100, 0); + EXPECT_EQ(b.remaining_rows(0, 0), 100); + EXPECT_EQ(b.remaining_rows(60, 9999), 40); + EXPECT_EQ(b.remaining_rows(100, 0), 0); + EXPECT_EQ(b.remaining_rows(200, 0), 0); +} + +TEST_F(BlockBudgetTest, RemainingRowsByteLimited) { + // max_rows=100, max_bytes=1000, current: 50 rows, 600 bytes + // avg = 12 bytes/row, byte_capacity = (1000-600)/12 = 33 + // row_capacity = 100 - 50 = 50 + // result = min(50, 33) = 33 + BlockBudget b(100, 1000); + EXPECT_EQ(b.remaining_rows(50, 600), 33); +} + +TEST_F(BlockBudgetTest, RemainingRowsAlreadyOverByteBudget) { + BlockBudget b(100, 1000); + EXPECT_EQ(b.remaining_rows(50, 1000), 0); + EXPECT_EQ(b.remaining_rows(50, 2000), 0); +} + +TEST_F(BlockBudgetTest, RemainingRowsZeroCurrentRows) { + // No rows yet → can't estimate avg_row_bytes, fall back to row capacity + BlockBudget b(100, 1000); + EXPECT_EQ(b.remaining_rows(0, 0), 100); +} + +TEST_F(BlockBudgetTest, RemainingRowsZeroCurrentBytes) { + // Has rows but zero bytes → can't estimate avg, fall back to row capacity + BlockBudget b(100, 1000); + EXPECT_EQ(b.remaining_rows(50, 0), 50); +} + +// ── has_byte_budget ───────────────────────────────────────────────────────── + +TEST_F(BlockBudgetTest, HasByteBudget) { + EXPECT_FALSE(BlockBudget(100, 0).has_byte_budget()); + EXPECT_TRUE(BlockBudget(100, 1).has_byte_budget()); +} + +} // namespace doris --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
