[feat](storage) Implement adaptive batch size for SegmentIterator

Introduce an EWMA-based AdaptiveBlockSizePredictor that dynamically
adjusts SegmentIterator chunk row counts so each output Block stays
near the session-variable-configured preferred_block_size_bytes target.
A complementary byte-budget stop condition is added to all BlockReader
and VCollectIterator accumulation loops so the final Block returned to
the upper layer is also bounded.

Key changes:
- New BE config: enable_adaptive_batch_size (default: true)
- New session variables: preferred_block_size_bytes (8 MB),
  preferred_max_column_in_block_size_bytes (1 MB)
- New Thrift fields 204/205 for FE→BE propagation
- AdaptiveBlockSizePredictor: EWMA per-row and per-column byte
  estimator with conservative segment-metadata bootstrap
- SegmentIterator: predicts rows before each next_batch call;
  updates EWMA on success path
- BlockReader: byte-stop in _replace_key_next_block,
  _unique_key_next_block, _agg_key_next_block
- VCollectIterator: byte-stop in Level1Iterator::_merge_next
- Unit tests and regression tests included

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: mrhhsg

be/src/common/config.cpp

@@ -433,6 +433,8 @@ DEFINE_mInt32(pk_index_page_cache_stale_sweep_time_sec, "600");
 DEFINE_mBool(enable_low_cardinality_optimize, "true");
 DEFINE_Bool(enable_low_cardinality_cache_code, "true");
 
+DEFINE_mBool(enable_adaptive_batch_size, "true");
+
 // be policy
 // whether check compaction checksum
 DEFINE_mBool(enable_compaction_checksum, "false");

be/src/common/config.h

@@ -498,6 +498,11 @@ DECLARE_mInt32(pk_index_page_cache_stale_sweep_time_sec);
 DECLARE_mBool(enable_low_cardinality_optimize);
 DECLARE_Bool(enable_low_cardinality_cache_code);
 
+// Adaptive batch size: dynamically adjust SegmentIterator chunk row count using EWMA
+// so that each output block stays close to preferred_block_size_bytes.
+// When false, the fixed batch_size row behaviour is preserved.
+DECLARE_mBool(enable_adaptive_batch_size);
+
 // be policy
 // whether check compaction checksum
 DECLARE_mBool(enable_compaction_checksum);

be/src/core/block/block.h

@@ -352,6 +352,15 @@ class Block {
     void clear_column_mem_not_keep(const std::vector<bool>& column_keep_flags,
                                    bool need_keep_first);
 
+    // Helper: sum byte_size() of all mutable columns.
+    static inline size_t columns_byte_size(const MutableColumns& cols) {
+        size_t total = 0;
+        for (const auto& col : cols) {
+            total += col->byte_size();
+        }
+        return total;
+    }
+
 private:
     void erase_impl(size_t position);
 };

be/src/exec/operator/olap_scan_operator.cpp

@@ -379,6 +379,11 @@ Status OlapScanLocalState::_init_profile() {
     _variant_doc_value_column_iter_count =
             ADD_COUNTER(_segment_profile, "VariantDocValueColumnIterCount", TUnit::UNIT);
 
+    _adaptive_batch_predict_min_rows_counter =
+            ADD_COUNTER(_segment_profile, "AdaptiveBatchPredictMinRows", TUnit::UNIT);
+    _adaptive_batch_predict_max_rows_counter =
+            ADD_COUNTER(_segment_profile, "AdaptiveBatchPredictMaxRows", TUnit::UNIT);
+
     return Status::OK();
 }
 

be/src/exec/operator/olap_scan_operator.h

@@ -309,6 +309,9 @@ class OlapScanLocalState final : public ScanLocalState<OlapScanLocalState> {
     // Variant subtree: times selecting doc snapshot all iterator (merge doc snapshot into root)
     RuntimeProfile::Counter* _variant_doc_value_column_iter_count = nullptr;
 
+    RuntimeProfile::Counter* _adaptive_batch_predict_min_rows_counter = nullptr;
+    RuntimeProfile::Counter* _adaptive_batch_predict_max_rows_counter = nullptr;
+
     std::vector<TabletWithVersion> _tablets;
     std::vector<TabletReadSource> _read_sources;
 

be/src/exec/scan/olap_scanner.cpp

@@ -164,6 +164,10 @@ Status OlapScanner::prepare() {
     // value (e.g. select a from t where a .. and b ... limit 1),
     // it will be very slow when reading data in segment iterator
     _tablet_reader->set_batch_size(_state->batch_size());
+    // Adaptive batch size: pass session variable values to the storage reader.
+    // The reader will forward them to RowsetReaderContext -> StorageReadOptions -> predictor.
+    _tablet_reader->set_preferred_block_size_bytes(_state->preferred_block_size_bytes());
+    _tablet_reader->set_preferred_max_col_bytes(_state->preferred_max_column_in_block_size_bytes());
     TabletSchemaSPtr cached_schema;
     std::string schema_key;
     {
@@ -840,6 +844,19 @@ void OlapScanner::_collect_profile_before_close() {
     COUNTER_UPDATE(local_state->_variant_doc_value_column_iter_count,
                    stats.variant_doc_value_column_iter_count);
 
+    if (stats.adaptive_batch_size_predict_max_rows > 0) {
+        auto cur_min = local_state->_adaptive_batch_predict_min_rows_counter->value();
+        if (cur_min == 0 || stats.adaptive_batch_size_predict_min_rows < cur_min) {
+            COUNTER_SET(local_state->_adaptive_batch_predict_min_rows_counter,
+                        stats.adaptive_batch_size_predict_min_rows);
+        }
+        auto cur_max = local_state->_adaptive_batch_predict_max_rows_counter->value();
+        if (stats.adaptive_batch_size_predict_max_rows > cur_max) {
+            COUNTER_SET(local_state->_adaptive_batch_predict_max_rows_counter,
+                        stats.adaptive_batch_size_predict_max_rows);
+        }
+    }
+
     InvertedIndexProfileReporter inverted_index_profile;
     inverted_index_profile.update(local_state->_index_filter_profile.get(),
                                   &stats.inverted_index_stats);

be/src/runtime/runtime_state.h

@@ -139,6 +139,20 @@ class RuntimeState {
     const DescriptorTbl& desc_tbl() const { return *_desc_tbl; }
     void set_desc_tbl(const DescriptorTbl* desc_tbl) { _desc_tbl = desc_tbl; }
     MOCK_FUNCTION int batch_size() const { return _query_options.batch_size; }
+    MOCK_FUNCTION size_t preferred_block_size_bytes() const {
+        if (_query_options.__isset.preferred_block_size_bytes) {
+            auto v = _query_options.preferred_block_size_bytes;
+            return v > 0 ? static_cast<size_t>(v) : 0;
+        }
+        return 8388608UL; // 8MB default
+    }
+    MOCK_FUNCTION size_t preferred_max_column_in_block_size_bytes() const {
+        if (_query_options.__isset.preferred_max_column_in_block_size_bytes) {
+            auto v = _query_options.preferred_max_column_in_block_size_bytes;
+            return v > 0 ? static_cast<size_t>(v) : 0;
+        }
+        return 1048576UL; // 1MB default
+    }
     int query_parallel_instance_num() const { return _query_options.parallel_instance; }
     int max_errors() const { return _query_options.max_errors; }
     int execution_timeout() const {

be/src/storage/iterator/block_reader.cpp

@@ -30,6 +30,7 @@
 // IWYU pragma: no_include <opentelemetry/common/threadlocal.h>
 #include "cloud/config.h"
 #include "common/compiler_util.h" // IWYU pragma: keep
+#include "common/config.h"
 #include "common/status.h"
 #include "core/block/column_with_type_and_name.h"
 #include "core/column/column_nullable.h"
@@ -386,6 +387,16 @@ Status BlockReader::_replace_key_next_block(Block* block, bool* eof) {
                 break;
             }
         }
+        // Byte-budget check: after the inner loop _next_row is either EOF or the next different
+        // key, so it is safe to stop accumulating here without repeating any row.
+        if (config::enable_adaptive_batch_size && _reader_context.preferred_block_size_bytes > 0 &&
+            Block::columns_byte_size(target_columns) >=
+                    _reader_context.preferred_block_size_bytes) {
+            if (UNLIKELY(_reader_context.record_rowids)) {
+                _block_row_locations.resize(target_block_row);
+            }
+            break;
+        }
     }
     _merged_rows += merged_row;
     return Status::OK();
@@ -480,6 +491,14 @@ Status BlockReader::_agg_key_next_block(Block* block, bool* eof) {
             if (target_block_row == _reader_context.batch_size) {
                 break;
             }
+            // Byte-budget check at group boundary: _next_row is the first row of the new group
+            // and is still pending (not yet inserted), so stopping here is safe.
+            if (config::enable_adaptive_batch_size &&
+                _reader_context.preferred_block_size_bytes > 0 &&
+                Block::columns_byte_size(target_columns) >=
+                        _reader_context.preferred_block_size_bytes) {
+                break;
+            }
             _agg_data_counters.push_back(_last_agg_data_counter);
             _last_agg_data_counter = 0;
 
@@ -539,6 +558,15 @@ Status BlockReader::_unique_key_next_block(Block* block, bool* eof) {
             LOG(WARNING) << "next failed: " << res;
             return res;
         }
+        // Byte-budget check: _next_row is already saved so stopping here is safe.
+        if (config::enable_adaptive_batch_size && _reader_context.preferred_block_size_bytes > 0 &&
+            Block::columns_byte_size(target_columns) >=
+                    _reader_context.preferred_block_size_bytes) {
+            if (UNLIKELY(_reader_context.record_rowids)) {
+                _block_row_locations.resize(target_block_row);
+            }
+            break;
+        }
     } while (target_block_row < _reader_context.batch_size);
 
     if (_delete_sign_available) {

be/src/storage/iterator/block_reader.h

@@ -24,6 +24,7 @@
 #include <utility>
 #include <vector>
 
+#include "common/config.h"
 #include "common/status.h"
 #include "core/block/block.h"
 #include "core/column/column.h"
@@ -54,6 +55,11 @@ class BlockReader final : public TabletReader {
         return _vcollect_iter.update_profile(profile);
     }
 
+    // Returns the configured preferred output block byte budget; 0 when adaptive is disabled.
+    size_t preferred_block_size_bytes() const override {
+        return config::enable_adaptive_batch_size ? _reader_context.preferred_block_size_bytes : 0;
+    }
+
 private:
     // Directly read row from rowset and pass to upper caller. No need to do aggregation.
     // This is usually used for DUPLICATE KEY tables

be/src/storage/iterator/vcollect_iterator.cpp

@@ -29,6 +29,7 @@
 
 #include "common/cast_set.h"
 #include "common/compiler_util.h" // IWYU pragma: keep
+#include "common/config.h"
 #include "common/status.h"
 #include "core/block/column_with_type_and_name.h"
 #include "core/column/column.h"
@@ -857,6 +858,24 @@ Status VCollectIterator::Level1Iterator::_merge_next(Block* block) {
             block->set_columns(std::move(target_columns));
             return Status::OK();
         }
+        // Byte-budget check: _merge_next() has already advanced _ref to the next unread row,
+        // so it is safe to stop here without duplicating any data.
+        if (_reader->preferred_block_size_bytes() > 0 &&
+            Block::columns_byte_size(target_columns) >= _reader->preferred_block_size_bytes()) {
+            if (continuous_row_in_block > 0) {
+                const auto& src_block = pre_row_ref.block;
+                for (size_t i = 0; i < column_count; ++i) {
+                    target_columns[i]->insert_range_from(*(src_block->get_by_position(i).column),
+                                                         pre_row_ref.row_pos,
+                                                         continuous_row_in_block);
+                }
+            }
+            if (UNLIKELY(_reader->_reader_context.record_rowids)) {
+                _block_row_locations.resize(target_block_row);
+            }
+            block->set_columns(std::move(target_columns));
+            return Status::OK();
+        }
         if (continuous_row_in_block == 0) {
             pre_row_ref = _ref;
             continue;