[Refactor] Extract some common code logic in TabletSchema (StarRocks#44193)

trueeyu · leiyang0324 · commit b1421b5c2d9a · 2024-04-19T14:46:12.000+08:00
Signed-off-by: trueeyu &lt;lxhhust350@qq.com&gt;
diff --git a/be/src/exec/pipeline/scan/olap_chunk_source.cpp b/be/src/exec/pipeline/scan/olap_chunk_source.cpp
@@ -412,18 +412,13 @@ Status OlapChunkSource::_init_olap_reader(RuntimeState* runtime_state) {
 
     auto scope = IOProfiler::scope(IOProfiler::TAG_QUERY, _scan_range->tablet_id);
 
-    auto tablet_schema_ptr = _tablet->tablet_schema();
-    _tablet_schema = TabletSchema::copy(tablet_schema_ptr);
-
     // if column_desc come from fe, reset tablet schema
     if (_scan_node->thrift_olap_scan_node().__isset.columns_desc &&
         !_scan_node->thrift_olap_scan_node().columns_desc.empty() &&
         _scan_node->thrift_olap_scan_node().columns_desc[0].col_unique_id >= 0) {
-        _tablet_schema->clear_columns();
-        for (const auto& column_desc : _scan_node->thrift_olap_scan_node().columns_desc) {
-            _tablet_schema->append_column(TabletColumn(column_desc));
-        }
-        _tablet_schema->generate_sort_key_idxes();
+        _tablet_schema = TabletSchema::copy(_tablet->tablet_schema(), _scan_node->thrift_olap_scan_node().columns_desc);
+    } else {
+        _tablet_schema = _tablet->tablet_schema();
     }
 
     RETURN_IF_ERROR(_init_global_dicts(&_params));
diff --git a/be/src/exec/pipeline/scan/olap_chunk_source.h b/be/src/exec/pipeline/scan/olap_chunk_source.h
@@ -85,7 +85,7 @@ class OlapChunkSource final : public ChunkSource {
 
     ObjectPool _obj_pool;
     TabletSharedPtr _tablet;
-    std::shared_ptr<TabletSchema> _tablet_schema;
+    TabletSchemaCSPtr _tablet_schema;
     int64_t _version = 0;
 
     RuntimeState* _runtime_state = nullptr;
diff --git a/be/src/exec/tablet_scanner.cpp b/be/src/exec/tablet_scanner.cpp
@@ -49,17 +49,12 @@ Status TabletScanner::init(RuntimeState* runtime_state, const TabletScannerParam
     RETURN_IF_ERROR(Expr::clone_if_not_exists(runtime_state, &_pool, *params.conjunct_ctxs, &_conjunct_ctxs));
     RETURN_IF_ERROR(_get_tablet(params.scan_range));
 
-    auto tablet_schema_ptr = _tablet->tablet_schema();
-    _tablet_schema = TabletSchema::copy(tablet_schema_ptr);
-
     // if column_desc come from fe, reset tablet schema
     if (_parent->_olap_scan_node.__isset.columns_desc && !_parent->_olap_scan_node.columns_desc.empty() &&
         _parent->_olap_scan_node.columns_desc[0].col_unique_id >= 0) {
-        _tablet_schema->clear_columns();
-        for (const auto& column_desc : _parent->_olap_scan_node.columns_desc) {
-            _tablet_schema->append_column(TabletColumn(column_desc));
-        }
-        _tablet_schema->generate_sort_key_idxes();
+        _tablet_schema = TabletSchema::copy(_tablet->tablet_schema(), _parent->_olap_scan_node.columns_desc);
+    } else {
+        _tablet_schema = _tablet->tablet_schema();
     }
 
     RETURN_IF_ERROR(_init_unused_output_columns(*params.unused_output_columns));
diff --git a/be/src/exec/tablet_scanner.h b/be/src/exec/tablet_scanner.h
@@ -105,7 +105,7 @@ class TabletScanner {
     std::shared_ptr<TabletReader> _reader;
 
     TabletSharedPtr _tablet;
-    TabletSchemaSPtr _tablet_schema;
+    TabletSchemaCSPtr _tablet_schema;
     int64_t _version = 0;
 
     // output columns of `this` TabletScanner, i.e, the final output columns of `get_chunk`.
diff --git a/be/src/storage/compaction.cpp b/be/src/storage/compaction.cpp
@@ -158,8 +158,7 @@ Status Compaction::_merge_rowsets_horizontally(size_t segment_iterator_num, Stat
                                                const TabletSchemaCSPtr& tablet_schema) {
     TRACE_COUNTER_SCOPE_LATENCY_US("merge_rowsets_latency_us");
     Schema schema = ChunkHelper::convert_schema(tablet_schema);
-    auto merge_tablet_schema = std::shared_ptr<TabletSchema>(TabletSchema::copy(tablet_schema));
-    TabletReader reader(_tablet, _output_rs_writer->version(), merge_tablet_schema, schema);
+    TabletReader reader(_tablet, _output_rs_writer->version(), tablet_schema, schema);
     TabletReaderParams reader_params;
     reader_params.reader_type = compaction_type();
     reader_params.profile = _runtime_profile.create_child("merge_rowsets");
diff --git a/be/src/storage/push_handler.cpp b/be/src/storage/push_handler.cpp
@@ -102,14 +102,12 @@ Status PushHandler::_do_streaming_ingestion(TabletSharedPtr tablet, const TPushR
             DeletePredicatePB del_pred;
             DeleteConditionHandler del_cond_handler;
             tablet_var.tablet->obtain_header_rdlock();
-            auto tablet_schema = TabletSchema::copy(tablet_var.tablet->tablet_schema());
+            TabletSchemaCSPtr tablet_schema;
             if (request.__isset.columns_desc && !request.columns_desc.empty() &&
                 request.columns_desc[0].col_unique_id >= 0) {
-                tablet_schema->clear_columns();
-                for (const auto& column_desc : request.columns_desc) {
-                    tablet_schema->append_column(TabletColumn(column_desc));
-                }
-                tablet_schema->generate_sort_key_idxes();
+                tablet_schema = TabletSchema::copy(tablet_var.tablet->tablet_schema(), request.columns_desc);
+            } else {
+                tablet_schema = tablet_var.tablet->tablet_schema();
             }
             res = del_cond_handler.generate_delete_predicate(*tablet_schema, request.delete_conditions, &del_pred);
             del_preds.push(del_pred);
@@ -122,13 +120,11 @@ Status PushHandler::_do_streaming_ingestion(TabletSharedPtr tablet, const TPushR
         }
     }
 
-    auto tablet_schema = std::shared_ptr<TabletSchema>(TabletSchema::copy(tablet_vars->at(0).tablet->tablet_schema()));
+    TabletSchemaCSPtr tablet_schema;
     if (request.__isset.columns_desc && !request.columns_desc.empty() && request.columns_desc[0].col_unique_id >= 0) {
-        tablet_schema->clear_columns();
-        for (const auto& column_desc : request.columns_desc) {
-            tablet_schema->append_column(TabletColumn(column_desc));
-        }
-        tablet_schema->generate_sort_key_idxes();
+        tablet_schema = TabletSchema::copy(tablet_vars->at(0).tablet->tablet_schema(), request.columns_desc);
+    } else {
+        tablet_schema = tablet_vars->at(0).tablet->tablet_schema();
     }
 
     Status st = Status::OK();
diff --git a/be/src/storage/rowset/rowset_meta.h b/be/src/storage/rowset/rowset_meta.h
@@ -245,18 +245,15 @@ class RowsetMeta {
         }
     }
 
-    void get_tablet_schema_pb(TabletSchemaPB* tablet_schema_pb) {
-        DCHECK(_schema != nullptr);
-        _schema->to_schema_pb(tablet_schema_pb);
-    }
-
     void set_tablet_schema(const TabletSchemaCSPtr& tablet_schema_ptr) {
         _rowset_meta_pb->clear_tablet_schema();
         TabletSchemaPB ts_pb;
         tablet_schema_ptr->to_schema_pb(&ts_pb);
         if (ts_pb.has_id() && ts_pb.id() != TabletSchema::invalid_id()) {
             _schema = GlobalTabletSchemaMap::Instance()->emplace(ts_pb).first;
         } else {
+            // Only for compatible, in very old versions, there is no schema id.
+            // If you fill with the default value, you cannot judge whether it is the same schema through the schema id.
             _schema = TabletSchemaCSPtr(TabletSchema::copy(tablet_schema_ptr));
         }
         _has_tablet_schema_pb = true;
diff --git a/be/src/storage/tablet_reader.cpp b/be/src/storage/tablet_reader.cpp
@@ -51,7 +51,7 @@ TabletReader::TabletReader(TabletSharedPtr tablet, const Version& version, Schem
 }
 
 TabletReader::TabletReader(TabletSharedPtr tablet, const Version& version, Schema schema,
-                           std::vector<RowsetSharedPtr> captured_rowsets, const TabletSchemaSPtr* tablet_schema)
+                           std::vector<RowsetSharedPtr> captured_rowsets, const TabletSchemaCSPtr* tablet_schema)
         : ChunkIterator(std::move(schema)),
           _tablet(std::move(tablet)),
           _version(version),
@@ -73,7 +73,7 @@ TabletReader::TabletReader(TabletSharedPtr tablet, const Version& version, Schem
     _tablet_schema = !tablet_schema ? _tablet->tablet_schema() : tablet_schema;
 }
 
-TabletReader::TabletReader(TabletSharedPtr tablet, const Version& version, const TabletSchemaSPtr& tablet_schema,
+TabletReader::TabletReader(TabletSharedPtr tablet, const Version& version, const TabletSchemaCSPtr& tablet_schema,
                            Schema schema)
         : ChunkIterator(std::move(schema)), _tablet(std::move(tablet)), _version(version) {
     _tablet_schema = tablet_schema;
diff --git a/be/src/storage/tablet_reader.h b/be/src/storage/tablet_reader.h
@@ -38,10 +38,10 @@ class TabletReader final : public ChunkIterator {
                  const TabletSchemaCSPtr& tablet_schema = nullptr);
     // *captured_rowsets* is captured forward before creating TabletReader.
     TabletReader(TabletSharedPtr tablet, const Version& version, Schema schema,
-                 std::vector<RowsetSharedPtr> captured_rowsets, const TabletSchemaSPtr* tablet_schema = nullptr);
+                 std::vector<RowsetSharedPtr> captured_rowsets, const TabletSchemaCSPtr* tablet_schema = nullptr);
     TabletReader(TabletSharedPtr tablet, const Version& version, Schema schema, bool is_key,
                  RowSourceMaskBuffer* mask_buffer, const TabletSchemaCSPtr& tablet_schema = nullptr);
-    TabletReader(TabletSharedPtr tablet, const Version& version, const TabletSchemaSPtr& tablet_schema, Schema schema);
+    TabletReader(TabletSharedPtr tablet, const Version& version, const TabletSchemaCSPtr& tablet_schema, Schema schema);
     ~TabletReader() override { close(); }
 
     void set_is_asc_hint(bool is_asc) { _is_asc_hint = is_asc; }
diff --git a/be/src/storage/tablet_schema.cpp b/be/src/storage/tablet_schema.cpp
@@ -375,6 +375,17 @@ std::unique_ptr<TabletSchema> TabletSchema::copy(const std::shared_ptr<const Tab
     return t_ptr;
 }
 
+TabletSchemaCSPtr TabletSchema::copy(const TabletSchemaCSPtr& src_schema, const std::vector<TColumn>& cols) {
+    auto dst_schema = std::make_unique<TabletSchema>();
+    dst_schema->copy_from(src_schema);
+    dst_schema->clear_columns();
+    for (const auto& col : cols) {
+        dst_schema->append_column(TabletColumn(col));
+    }
+    dst_schema->generate_sort_key_idxes();
+    return dst_schema;
+}
+
 void TabletSchema::_fill_index_map(const TabletIndex& index) {
     const auto idx_type = index.index_type();
     if (_index_map_col_unique_id.count(idx_type) <= 0) {
diff --git a/be/src/storage/tablet_schema.h b/be/src/storage/tablet_schema.h
@@ -237,15 +237,16 @@ class TabletIndex;
 class TabletSchema {
 public:
     using SchemaId = int64_t;
+    using TabletSchemaSPtr = std::shared_ptr<TabletSchema>;
     using TabletSchemaCSPtr = std::shared_ptr<const TabletSchema>;
 
-    static std::shared_ptr<TabletSchema> create(const TabletSchemaPB& schema_pb);
-    static std::shared_ptr<TabletSchema> create(const TabletSchemaPB& schema_pb, TabletSchemaMap* schema_map);
-    static std::shared_ptr<TabletSchema> create(const TabletSchemaCSPtr& tablet_schema,
-                                                const std::vector<int32_t>& column_indexes);
-    static std::shared_ptr<TabletSchema> create_with_uid(const TabletSchemaCSPtr& tablet_schema,
-                                                         const std::vector<uint32_t>& unique_column_ids);
-    static std::unique_ptr<TabletSchema> copy(const std::shared_ptr<const TabletSchema>& tablet_schema);
+    static TabletSchemaSPtr create(const TabletSchemaPB& schema_pb);
+    static TabletSchemaSPtr create(const TabletSchemaPB& schema_pb, TabletSchemaMap* schema_map);
+    static TabletSchemaSPtr create(const TabletSchemaCSPtr& tablet_schema, const std::vector<int32_t>& column_indexes);
+    static TabletSchemaSPtr create_with_uid(const TabletSchemaCSPtr& tablet_schema,
+                                            const std::vector<uint32_t>& unique_column_ids);
+    static std::unique_ptr<TabletSchema> copy(const TabletSchemaCSPtr& tablet_schema);
+    static TabletSchemaCSPtr copy(const TabletSchemaCSPtr& src_schema, const std::vector<TColumn>& cols);
 
     // Must be consistent with MaterializedIndexMeta.INVALID_SCHEMA_ID defined in
     // file ./fe/fe-core/src/main/java/com/starrocks/catalog/MaterializedIndexMeta.java
