1

Author: wyxxxcat

be/src/olap/rowset/segment_v2/segment_writer.cpp

@@ -792,11 +792,13 @@ Status SegmentWriter::append_block(const vectorized::Block* block, size_t row_po
             short_key_pos.push_back(_short_key_row_pos - _num_rows_written);
         }
     }
+    int64_t total_data_size = 0;
 
     // convert column data from engine format to storage layer format
     std::vector<vectorized::IOlapColumnDataAccessor*> key_columns;
     vectorized::IOlapColumnDataAccessor* seq_column = nullptr;
     for (size_t id = 0; id < _column_writers.size(); ++id) {
+        int64_t column_data_size = 0;
         // olap data convertor alway start from id = 0
         auto converted_result = _olap_data_convertor->convert_column_data(id);
         if (!converted_result.first.ok()) {
@@ -811,7 +813,18 @@ Status SegmentWriter::append_block(const vectorized::Block* block, size_t row_po
         }
         RETURN_IF_ERROR(_column_writers[id]->append(converted_result.second->get_nullmap(),
                                                     converted_result.second->get_data(), num_rows));
+
+        // estimate column data size for flush memtable, may be inaccurate at low cardinality
+        column_data_size += _column_writers[id]->estimate_buffer_size();
+        total_data_size += column_data_size;
+        auto origin_column_data_size = _footer.columns(id).total_data_size();
+        _footer.mutable_columns(id)->set_total_data_size(origin_column_data_size +
+                                                         column_data_size);
     }
+
+    auto origin_data_footprint = _footer.data_footprint();
+    _footer.set_data_footprint(origin_data_footprint + total_data_size);
+
     if (_has_key) {
         if (_is_mow_with_cluster_key()) {
             // for now we don't need to query short key index for CLUSTER BY feature,

be/src/olap/rowset/segment_v2/vertical_segment_writer.cpp

@@ -54,6 +54,7 @@
 #include "olap/short_key_index.h"
 #include "olap/tablet_schema.h"
 #include "olap/utils.h"
+#include "runtime/define_primitive_type.h"
 #include "runtime/exec_env.h"
 #include "runtime/memory/mem_tracker.h"
 #include "service/point_query_executor.h"
@@ -62,6 +63,7 @@
 #include "util/debug_points.h"
 #include "util/faststring.h"
 #include "util/key_util.h"
+#include "vec/columns/column_map.h"
 #include "vec/columns/column_nullable.h"
 #include "vec/columns/column_vector.h"
 #include "vec/columns/columns_number.h"
@@ -1174,7 +1176,9 @@ Status VerticalSegmentWriter::write_batch() {
     vectorized::IOlapColumnDataAccessor* seq_column = nullptr;
     // the key is cluster key column unique id
     std::map<uint32_t, vectorized::IOlapColumnDataAccessor*> cid_to_column;
+    int64_t total_data_size = 0;
     for (uint32_t cid = 0; cid < _tablet_schema->num_columns(); ++cid) {
+        int64_t column_data_size = 0;
         RETURN_IF_ERROR(_create_column_writer(cid, _tablet_schema->column(cid), _tablet_schema));
         for (auto& data : _batched_blocks) {
             RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_columns(
@@ -1200,17 +1204,28 @@ Status VerticalSegmentWriter::write_batch() {
             }
             RETURN_IF_ERROR(_column_writers[cid]->append(column->get_nullmap(), column->get_data(),
                                                          data.num_rows));
+            if (_data_dir != nullptr &&
+                _data_dir->reach_capacity_limit(_column_writers[cid]->estimate_buffer_size())) {
+                return Status::Error<DISK_REACH_CAPACITY_LIMIT>("disk {} exceed capacity limit.",
+                                                                _data_dir->path_hash());
+            }
+
+            // estimate column data size for flush memtable, may be inaccurate at low cardinality
+            column_data_size += _column_writers[cid]->estimate_buffer_size();
+            total_data_size += column_data_size;
+            auto origin_data_size = _footer.columns(cid).total_data_size();
+            _footer.mutable_columns(cid)->set_total_data_size(origin_data_size + column_data_size);
+
+            RETURN_IF_ERROR(_column_writers[cid]->finish());
+            RETURN_IF_ERROR(_column_writers[cid]->write_data());
+
             _olap_data_convertor->clear_source_content();
         }
-        if (_data_dir != nullptr &&
-            _data_dir->reach_capacity_limit(_column_writers[cid]->estimate_buffer_size())) {
-            return Status::Error<DISK_REACH_CAPACITY_LIMIT>("disk {} exceed capacity limit.",
-                                                            _data_dir->path_hash());
-        }
-        RETURN_IF_ERROR(_column_writers[cid]->finish());
-        RETURN_IF_ERROR(_column_writers[cid]->write_data());
     }
 
+    auto origin_data_footprint = _footer.data_footprint();
+    _footer.set_data_footprint(origin_data_footprint + total_data_size);
+
     for (auto& data : _batched_blocks) {
         _olap_data_convertor->set_source_content(data.block, data.row_pos, data.num_rows);
         RETURN_IF_ERROR(_generate_key_index(data, key_columns, seq_column, cid_to_column));

gensrc/proto/segment_v2.proto

@@ -197,6 +197,8 @@ message ColumnMetaPB {
     optional bool result_is_nullable = 18; // used on agg_state type
     optional string function_name = 19; // used on agg_state type
     optional int32 be_exec_version = 20; // used on agg_state type
+
+    optional uint64 total_data_size = 21;
 }
 
 message PrimaryKeyIndexMetaPB {