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 c64212666ccc055ad73b1d9f7da6f9887a88fbdc Author: Hu Shenggang <[email protected]> AuthorDate: Tue Apr 7 18:30:54 2026 +0800 [fix](be) Enforce byte budget in distinct streaming agg and topN path ### What problem does this PR solve? Issue Number: close #xxx Related PR: #xxx Problem Summary: Two code paths lacked byte-budget enforcement for the adaptive batch size feature: 1. **Distinct Streaming Aggregation**: The split logic in _distinct_pre_agg_with_serialized_key only enforced row-count cap (batch_size). For wide grouping keys, this could produce blocks exceeding preferred_block_size_bytes within a single push(). Fix: Added byte-budget-aware split — estimates remaining byte budget using average row size and caps max_rows_to_add accordingly. Excess rows go to _cache_block for the next iteration. 2. **TopN storage path** (_topn_next): The one-shot topN computation returned the entire result in a single block without consulting byte budget. For large LIMIT values on wide tables, this could produce blocks far exceeding preferred_block_size_bytes. Fix: After computing the full topN result, if it exceeds the byte budget, the result is stored and returned in chunks. Each chunk respects both batch_max_rows and the byte budget. When byte budget is disabled, behavior is unchanged (zero overhead). ### Release note None ### Check List (For Author) - Test: Unit Test (78 + 105 = 183 tests passed including DistinctStreamingAgg, Sorter, Operator, and adaptive batch size tests) - Behavior changed: No (feature is gated on preferred_block_size_bytes > 0) - Does this need documentation: No Co-authored-by: Copilot <[email protected]> --- .../distinct_streaming_aggregation_operator.cpp | 28 +++++++++--- .../distinct_streaming_aggregation_operator.h | 1 + be/src/storage/iterator/vcollect_iterator.cpp | 51 ++++++++++++++++++++++ be/src/storage/iterator/vcollect_iterator.h | 5 +++ 4 files changed, 79 insertions(+), 6 deletions(-) diff --git a/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp b/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp index f8dcb2aa17a..64aa5f1acc0 100644 --- a/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp +++ b/be/src/exec/operator/distinct_streaming_aggregation_operator.cpp @@ -38,6 +38,7 @@ DistinctStreamingAggLocalState::DistinctStreamingAggLocalState(RuntimeState* sta OperatorXBase* parent) : PipelineXLocalState<FakeSharedState>(state, parent), batch_size(state->block_max_rows()), + block_max_bytes(state->block_max_bytes()), _agg_data(std::make_unique<DistinctDataVariants>()), _child_block(Block::create_unique()), _aggregated_block(Block::create_unique()), @@ -207,16 +208,31 @@ Status DistinctStreamingAggLocalState::_distinct_pre_agg_with_serialized_key( } } else { DCHECK_EQ(_cache_block.rows(), 0); - // is output row > batch_size, split some to cache_block - if (out_block->rows() + _distinct_row.size() > batch_size) { - size_t split_size = batch_size - out_block->rows(); + // 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); + } + } + } + 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(); key_columns[i]->append_data_by_selector(output_dst, _distinct_row, 0, - split_size); + max_rows_to_add); auto cache_dst = _cache_block.get_by_position(i).column->assume_mutable(); - key_columns[i]->append_data_by_selector(cache_dst, _distinct_row, split_size, - _distinct_row.size()); + key_columns[i]->append_data_by_selector(cache_dst, _distinct_row, + max_rows_to_add, _distinct_row.size()); } } else { for (int i = 0; i < key_size; ++i) { diff --git a/be/src/exec/operator/distinct_streaming_aggregation_operator.h b/be/src/exec/operator/distinct_streaming_aggregation_operator.h index 07e015af86e..9bc21acf429 100644 --- a/be/src/exec/operator/distinct_streaming_aggregation_operator.h +++ b/be/src/exec/operator/distinct_streaming_aggregation_operator.h @@ -69,6 +69,7 @@ private: bool _should_expand_hash_table = true; bool _stop_emplace_flag = false; const int batch_size; + const size_t block_max_bytes; std::unique_ptr<DistinctDataVariants> _agg_data = nullptr; // group by k1,k2 VExprContextSPtrs _probe_expr_ctxs; diff --git a/be/src/storage/iterator/vcollect_iterator.cpp b/be/src/storage/iterator/vcollect_iterator.cpp index b8f54b91dbe..dd41adf98fc 100644 --- a/be/src/storage/iterator/vcollect_iterator.cpp +++ b/be/src/storage/iterator/vcollect_iterator.cpp @@ -304,6 +304,11 @@ Status VCollectIterator::_topn_next(Block* block) { return Status::Error<END_OF_FILE>(""); } + // If we already computed the full topN result, return the next chunk. + if (_topn_result_block.rows() > 0) { + return _topn_next_chunk(block); + } + auto clone_block = block->clone_empty(); /* select id, "${tR2}", @@ -492,10 +497,56 @@ Status VCollectIterator::_topn_next(Block* block) { << " mutable_block.rows()=" << mutable_block.rows(); *block = mutable_block.to_block(); + // If byte budget is active and the result exceeds it, split into chunks. + size_t preferred_bytes = _reader->preferred_block_size_bytes(); + if (preferred_bytes > 0 && block->rows() > 1 && block->bytes() > preferred_bytes) { + _topn_result_block = std::move(*block); + _topn_result_offset = 0; + return _topn_next_chunk(block); + } + _topn_eof = true; return block->rows() > 0 ? Status::OK() : Status::Error<END_OF_FILE>(""); } +Status VCollectIterator::_topn_next_chunk(Block* block) { + size_t total_rows = _topn_result_block.rows(); + if (_topn_result_offset >= total_rows) { + _topn_result_block.clear(); + _topn_eof = true; + return Status::Error<END_OF_FILE>(""); + } + + size_t remaining = total_rows - _topn_result_offset; + size_t chunk_rows = remaining; + + // Cap by batch_max_rows. + chunk_rows = std::min(chunk_rows, (size_t)_reader->batch_max_rows()); + + // Cap by byte budget. + size_t preferred_bytes = _reader->preferred_block_size_bytes(); + if (preferred_bytes > 0 && total_rows > 0) { + size_t avg_row_bytes = std::max<size_t>(1, _topn_result_block.bytes() / total_rows); + chunk_rows = std::min(chunk_rows, std::max<size_t>(1, preferred_bytes / avg_row_bytes)); + } + + *block = _topn_result_block.clone_empty(); + auto cols = block->mutate_columns(); + for (size_t i = 0; i < cols.size(); ++i) { + cols[i]->insert_range_from(*_topn_result_block.get_by_position(i).column, + _topn_result_offset, chunk_rows); + } + block->set_columns(std::move(cols)); + + _topn_result_offset += chunk_rows; + if (_topn_result_offset >= total_rows) { + _topn_result_block.clear(); + _topn_eof = true; + } + + return Status::OK(); +} + bool VCollectIterator::BlockRowPosComparator::operator()(const size_t& lpos, const size_t& rpos) const { int ret = _mutable_block->compare_at(lpos, rpos, _compare_columns, *_mutable_block, -1); diff --git a/be/src/storage/iterator/vcollect_iterator.h b/be/src/storage/iterator/vcollect_iterator.h index 199f7823565..fc394065093 100644 --- a/be/src/storage/iterator/vcollect_iterator.h +++ b/be/src/storage/iterator/vcollect_iterator.h @@ -94,6 +94,8 @@ public: private: // next for topn query Status _topn_next(Block* block); + // Return next chunk from precomputed topN result. + Status _topn_next_chunk(Block* block); class BlockRowPosComparator { public: @@ -356,6 +358,9 @@ private: // for topn next size_t _topn_limit = 0; bool _topn_eof = false; + // For chunked topN output when result exceeds byte budget. + Block _topn_result_block; + size_t _topn_result_offset = 0; std::vector<RowSetSplits> _rs_splits; // General limit pushdown for DUP_KEYS and UNIQUE_KEYS with MOW (non-merge path). --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
