Pre-decode value in write cf

dbms/src/Raft/RaftService.cpp

@@ -59,7 +59,7 @@ RaftService::RaftService(const std::string & address_, DB::Context & db_context_
             region = it->second;
             regions_to_decode.erase(it);
         }
-        region->tryDecodeDefaultCF();
+        region->tryPreDecodeTiKVValue();
         return true;
     });
 

dbms/src/Storages/Transaction/ExtraCFData.h

@@ -6,6 +6,9 @@ namespace DB
 {
 
 struct RegionDefaultCFDataTrait;
+struct RegionWriteCFDataTrait;
+
+using ExtraCFDataQueue = std::deque<std::shared_ptr<const TiKVValue>>;
 
 template <typename Trait>
 struct ExtraCFData
@@ -16,7 +19,6 @@ struct ExtraCFData
 template <>
 struct ExtraCFData<RegionDefaultCFDataTrait>
 {
-    using Queue = std::deque<std::shared_ptr<const TiKVValue>>;
     mutable std::mutex default_cf_decode_mutex;
 
     ExtraCFData() = default;
@@ -27,13 +29,13 @@ struct ExtraCFData<RegionDefaultCFDataTrait>
         queue.push_back(e);
     }
 
-    std::optional<Queue> popAll()
+    std::optional<ExtraCFDataQueue> popAll()
     {
         std::lock_guard<std::mutex> lock(default_cf_decode_mutex);
         if (queue.empty())
             return {};
 
-        Queue res;
+        ExtraCFDataQueue res;
         queue.swap(res);
         return res;
     }
@@ -61,7 +63,19 @@ struct ExtraCFData<RegionDefaultCFDataTrait>
     }
 
 private:
-    Queue queue;
+    ExtraCFDataQueue queue;
+};
+
+template <>
+struct ExtraCFData<RegionWriteCFDataTrait> : ExtraCFData<RegionDefaultCFDataTrait>
+{
+    using Base = ExtraCFData<RegionDefaultCFDataTrait>;
+    void add(const std::shared_ptr<const TiKVValue> & e)
+    {
+        if (!e)
+            return;
+        Base::add(e);
+    }
 };
 
 } // namespace DB

dbms/src/Storages/Transaction/Region.cpp

@@ -655,16 +655,12 @@ void Region::doDeleteRange(const std::string & cf, const TiKVKey & start_key, co
 
 std::tuple<RegionVersion, RegionVersion, RegionRange> Region::dumpVersionRange() const { return meta.dumpVersionRange(); }
 
-void Region::tryDecodeDefaultCF()
+void tryPreDecodeTiKVValue(std::optional<ExtraCFDataQueue> && values)
 {
-    ExtraCFData<RegionDefaultCFDataTrait>::Queue values;
-    {
-        auto res = data.defaultCF().getExtra().popAll();
-        if (!res)
-            return;
-        values = std::move(*res);
-    }
-    for (const auto & val : values)
+    if (!values)
+        return;
+
+    for (const auto & val : *values)
     {
         auto & decoded_row_info = val->extraInfo();
         if (decoded_row_info.load())
@@ -674,4 +670,10 @@ void Region::tryDecodeDefaultCF()
     }
 }
 
+void Region::tryPreDecodeTiKVValue()
+{
+    DB::tryPreDecodeTiKVValue(data.defaultCF().getExtra().popAll());
+    DB::tryPreDecodeTiKVValue(data.writeCF().getExtra().popAll());
+}
+
 } // namespace DB

dbms/src/Storages/Transaction/Region.h

@@ -164,7 +164,7 @@ class Region : public std::enable_shared_from_this<Region>
 
     static ColumnFamilyType getCf(const std::string & cf);
 
-    void tryDecodeDefaultCF();
+    void tryPreDecodeTiKVValue();
 
 private:
     Region() = delete;

dbms/src/Storages/Transaction/RegionBlockReader.cpp

@@ -125,10 +125,13 @@ void setPKVersionDel(ColumnUInt8 & delmark_col,
     }
 }
 
+using ColumnIdToInfoIndexMap = google::dense_hash_map<ColumnID, UInt16>;
+using SchemaAllColumnIds = google::dense_hash_set<ColumnID>;
+
 /// DecodeRowSkip function will try to jump over unnecessary field.
-bool DecodeRowSkip(const TiKVValue & value, const google::dense_hash_set<ColumnID> & column_ids_to_read,
-    const google::dense_hash_set<ColumnID> & schema_all_column_ids, DecodedRow & additional_decoded_row,
-    std::vector<DecodedRow::const_iterator> & decoded_col_iter)
+bool DecodeRowSkip(const TiKVValue & value, const ColumnIdToInfoIndexMap & column_id_to_info_index,
+    const SchemaAllColumnIds & schema_all_column_ids, DecodedRow & additional_decoded_row,
+    std::vector<DecodedRow::const_iterator> & decoded_col_iter, const bool force_decode)
 {
     const String & raw_value = value.getStr();
     size_t cursor = 0;
@@ -141,12 +144,17 @@ bool DecodeRowSkip(const TiKVValue & value, const google::dense_hash_set<ColumnI
             break;
 
         ColumnID col_id = f.get<ColumnID>();
+
         column_cnt++;
-        if (!schema_all_column_ids.count(col_id))
+
+        if (schema_matches && !schema_all_column_ids.count(col_id))
         {
             schema_matches = false;
+            if (!force_decode)
+                return schema_matches;
         }
-        if (!column_ids_to_read.count(col_id))
+
+        if (!column_id_to_info_index.count(col_id))
         {
             SkipDatum(cursor, raw_value);
         }
@@ -156,6 +164,7 @@ bool DecodeRowSkip(const TiKVValue & value, const google::dense_hash_set<ColumnI
             decoded_col_iter.emplace_back(additional_decoded_row.cend() - 1);
         }
     }
+
     if (column_cnt != schema_all_column_ids.size())
     {
         schema_matches = false;
@@ -167,9 +176,9 @@ bool DecodeRowSkip(const TiKVValue & value, const google::dense_hash_set<ColumnI
 }
 
 /// DecodeRow function will try to get pre-decoded fields from value, if is none, just decode its str.
-bool DecodeRow(const TiKVValue & value, const google::dense_hash_set<ColumnID> & column_ids_to_read,
-    const google::dense_hash_set<ColumnID> & schema_all_column_ids, DecodedRow & additional_decoded_row,
-    std::vector<DecodedRow::const_iterator> & decoded_col_iter)
+bool DecodeRow(const TiKVValue & value, const ColumnIdToInfoIndexMap & column_id_to_info_index,
+    const SchemaAllColumnIds & schema_all_column_ids, DecodedRow & additional_decoded_row,
+    std::vector<DecodedRow::const_iterator> & decoded_col_iter, const bool force_decode)
 {
     auto & decoded_row_info = value.extraInfo();
     const DecodedRow * id_fields_ptr = decoded_row_info.load();
@@ -183,20 +192,28 @@ bool DecodeRow(const TiKVValue & value, const google::dense_hash_set<ColumnID> &
         {
             const auto & ele = *it;
             const auto & col_id = ele.col_id;
-            if (!schema_all_column_ids.count(col_id))
+
+            if (schema_matches && !schema_all_column_ids.count(col_id))
             {
                 schema_matches = false;
+                if (!force_decode)
+                    return schema_matches;
             }
-            if (column_ids_to_read.count(col_id))
+
+            if (column_id_to_info_index.count(col_id))
             {
                 decoded_col_iter.emplace_back(it);
             }
         }
+
+        if (id_fields.size() != schema_all_column_ids.size())
+            schema_matches = false;
+
         return schema_matches;
     }
     else
     {
-        return DecodeRowSkip(value, column_ids_to_read, schema_all_column_ids, additional_decoded_row, decoded_col_iter);
+        return DecodeRowSkip(value, column_id_to_info_index, schema_all_column_ids, additional_decoded_row, decoded_col_iter, force_decode);
     }
 }
 
@@ -219,20 +236,20 @@ std::tuple<Block, bool> readRegionBlock(const TiDB::TableInfo & table_info,
     google::dense_hash_map<ColumnID, std::shared_ptr<ColTypePair>> column_map;
     column_map.set_empty_key(EmptyColumnID);
 
-    google::dense_hash_map<ColumnID, size_t> column_id_to_info_index_map;
-    column_id_to_info_index_map.set_empty_key(EmptyColumnID);
+    ColumnIdToInfoIndexMap column_id_to_info_index;
+    column_id_to_info_index.set_empty_key(EmptyColumnID);
 
-    google::dense_hash_set<ColumnID> column_ids_to_read;
-    column_ids_to_read.set_empty_key(EmptyColumnID);
-
-    google::dense_hash_set<ColumnID> schema_all_column_ids;
+    SchemaAllColumnIds schema_all_column_ids;
     schema_all_column_ids.set_empty_key(EmptyColumnID);
 
+    if (table_info.columns.size() > std::numeric_limits<ColumnIdToInfoIndexMap::mapped_type>::max())
+        throw Exception("Too many columns in schema", ErrorCodes::LOGICAL_ERROR);
+
     for (size_t i = 0; i < table_info.columns.size(); i++)
     {
         auto & column_info = table_info.columns[i];
         ColumnID col_id = column_info.id;
-        String col_name = column_info.name;
+        const String & col_name = column_info.name;
         schema_all_column_ids.insert(col_id);
         if (std::find(column_names_to_read.begin(), column_names_to_read.end(), col_name) == column_names_to_read.end())
         {
@@ -244,12 +261,10 @@ std::tuple<Block, bool> readRegionBlock(const TiDB::TableInfo & table_info,
         if (table_info.pk_is_handle && column_info.hasPriKeyFlag())
             handle_col_id = col_id;
         else
-        {
-            column_ids_to_read.insert(col_id);
-            column_id_to_info_index_map.insert(std::make_pair(col_id, i));
-        }
+            column_id_to_info_index.insert(std::make_pair(col_id, i));
     }
-    if (column_names_to_read.size() - 3 != column_ids_to_read.size())
+
+    if (column_names_to_read.size() - 3 != column_id_to_info_index.size())
         throw Exception("schema doesn't contain needed columns.", ErrorCodes::LOGICAL_ERROR);
 
     if (!table_info.pk_is_handle)
@@ -315,42 +330,42 @@ std::tuple<Block, bool> readRegionBlock(const TiDB::TableInfo & table_info,
 
             if (write_type == Region::DelFlag)
             {
-                for (auto col_id : column_ids_to_read)
+                for (const auto & item : column_id_to_info_index)
                 {
-                    const auto & column = table_info.columns[column_id_to_info_index_map[col_id]];
+                    const auto & column = table_info.columns[item.second];
 
                     additional_decoded_row.emplace_back(column.id, GenDecodeRow(column));
                     decoded_col_iter.emplace_back(additional_decoded_row.cend() - 1);
                 }
             }
             else
             {
-                bool schema_matches
-                    = DecodeRow(*value_ptr, column_ids_to_read, schema_all_column_ids, additional_decoded_row, decoded_col_iter);
+                bool schema_matches = DecodeRow(
+                    *value_ptr, column_id_to_info_index, schema_all_column_ids, additional_decoded_row, decoded_col_iter, force_decode);
                 if (!schema_matches && !force_decode)
                     return std::make_tuple(block, false);
             }
 
             /// Modify `row` by adding missing column values or removing useless column values.
-            if (unlikely(decoded_col_iter.size() > column_ids_to_read.size()))
+            if (unlikely(decoded_col_iter.size() > column_id_to_info_index.size()))
             {
                 throw Exception("read unexpected columns.", ErrorCodes::LOGICAL_ERROR);
             }
 
             // redundant column values (column id not in current schema) has been dropped when decoding row
             // this branch handles the case when the row doesn't contain all the needed column
-            if (decoded_col_iter.size() < column_ids_to_read.size())
+            if (decoded_col_iter.size() < column_id_to_info_index.size())
             {
-                decoded_col_ids_set.clear();
+                decoded_col_ids_set.clear_no_resize();
                 for (const auto & e : decoded_col_iter)
                     decoded_col_ids_set.insert(e->col_id);
 
-                for (auto col_id : column_ids_to_read)
+                for (const auto & item : column_id_to_info_index)
                 {
-                    if (decoded_col_ids_set.count(col_id))
+                    if (decoded_col_ids_set.count(item.first))
                         continue;
 
-                    const auto & column = table_info.columns[column_id_to_info_index_map[col_id]];
+                    const auto & column = table_info.columns[item.second];
 
                     additional_decoded_row.emplace_back(column.id,
                         column.hasNoDefaultValueFlag() ? (column.hasNotNullFlag() ? GenDecodeRow(column) : Field())

dbms/src/Storages/Transaction/RegionCFDataBase.cpp

@@ -48,7 +48,10 @@ TableID RegionCFDataBase<Trait>::insert(const TableID table_id, std::pair<Key, V
     if (!ok)
         throw Exception("Found existing key in hex: " + getTiKVKey(kv_pair.second).toHex(), ErrorCodes::LOGICAL_ERROR);
 
-    extra.add(getTiKVValuePtr(it->second));
+    if constexpr (std::is_same_v<Trait, RegionWriteCFDataTrait>)
+        extra.add(Trait::getRowRawValuePtr(it->second));
+    else
+        extra.add(getTiKVValuePtr(it->second));
     return table_id;
 }
 

dbms/src/Storages/Transaction/RegionCFDataTrait.h

@@ -24,7 +24,7 @@ struct RegionWriteCFDataTrait
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>, DecodedWriteCFValue>;
     using Map = std::map<Key, Value>;
 
-    static std::pair<Key, Value> genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
+    static Map::value_type genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
     {
         HandleID handle_id = RecordKVFormat::getHandle(raw_key);
         Timestamp ts = RecordKVFormat::getTs(key);
@@ -34,6 +34,8 @@ struct RegionWriteCFDataTrait
                 std::move(decoded_val)}};
     }
 
+    static const std::shared_ptr<const TiKVValue> & getRowRawValuePtr(const Value & val) { return std::get<2>(std::get<2>(val)); }
+
     static UInt8 getWriteType(const Value & value) { return std::get<0>(std::get<2>(value)); }
 };
 
@@ -44,7 +46,7 @@ struct RegionDefaultCFDataTrait
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>>;
     using Map = std::map<Key, Value>;
 
-    static std::pair<Key, Value> genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
+    static Map::value_type genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
     {
         HandleID handle_id = RecordKVFormat::getHandle(raw_key);
         Timestamp ts = RecordKVFormat::getTs(key);
@@ -60,7 +62,7 @@ struct RegionLockCFDataTrait
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>, DecodedLockCFValue>;
     using Map = std::unordered_map<Key, Value>;
 
-    static std::pair<Key, Value> genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
+    static Map::value_type genKVPair(TiKVKey && key, const String & raw_key, TiKVValue && value)
     {
         HandleID handle_id = RecordKVFormat::getHandle(raw_key);
         auto decoded_val = RecordKVFormat::decodeLockCfValue(value);