Bugfix/issue 370 fix backfill with virtual column (#372)

* fix backfill from historical store with virtual column
* fix incorrectly caching shared records in tail cache

Author: yl-lisen

src/Core/Block.cpp

@@ -978,6 +978,30 @@ void Block::reorderColumnsInplace(const Block & header)
     // }
 }
 
+void Block::reorderColumnsInplace(const Names & names)
+{
+    assert(names.size() >= 1);
+
+    auto num_columns = columns();
+    if (num_columns <= 1)
+        return;
+
+    assert(num_columns >= names.size());
+
+    Block result;
+    result.reserve(names.size());
+
+    for (const auto & name : names)
+    {
+        auto & target_col = getByName(name);
+        result.insert(std::move(target_col));
+    }
+
+    /// Don't swap block.info
+    data.swap(result.data);
+    index_by_name.swap(result.index_by_name);
+}
+
 void Block::concat(const Block & other)
 {
     assert(blocksHaveEqualStructure(*this, other));

src/Core/Block.h

@@ -167,6 +167,9 @@ class Block
     /// in-place sort columns according to header
     void reorderColumnsInplace(const Block & header);
 
+    /// in-place sort columns according to names
+    void reorderColumnsInplace(const Names & names);
+
     void renameColumn(String new_name, size_t column_pos);
 
     /// Deep clone, use cautiously. Most of time, we don't need deepClone

src/NativeLog/Cache/TailCache.cpp

@@ -32,7 +32,12 @@ void TailCache::put(const StreamShard & stream_shard, RecordPtr record)
     {
         assert(entry);
         std::unique_lock lock(entry->mutex);
-        entry->shard_records.push_back(std::move(record));
+        if (record.use_count() == 1)
+            entry->shard_records.push_back(std::move(record));
+        else
+            /// Cache a copy record to prevent other references modified it
+            entry->shard_records.push_back(record->clone());
+
         if (entry->shard_records.size() > static_cast<size_t>(config.max_cached_entries_per_shard))
         {
             size -= entry->shard_records.front()->totalSerializedBytes();

src/NativeLog/Record/Record.h

@@ -58,6 +58,11 @@ struct Record
         setCodec(DB::CompressionMethodByte::NONE);
     }
 
+    RecordPtr clone() const
+    {
+        return std::make_shared<Record>(*this);
+    }
+
     RecordPtr clone(DB::Block && block_) const
     {
         auto r{std::make_shared<Record>(sn, std::move(block_))};

src/Processors/QueryPlan/ReadFromMergeTree.cpp

@@ -93,13 +93,15 @@ ReadFromMergeTree::ReadFromMergeTree(
     Poco::Logger * log_,
     MergeTreeDataSelectAnalysisResultPtr analyzed_result_ptr_,
     bool enable_parallel_reading,
-    std::function<std::shared_ptr<ISource>(Int64 &)> create_streaming_source_)
+    std::function<std::shared_ptr<ISource>(Int64 &)> create_streaming_source_,
+    const Names & column_names_to_return)
     : ISourceStep(DataStream{
         .header = MergeTreeBaseSelectProcessor::transformHeader(
             storage_snapshot_->getSampleBlockForColumns(real_column_names_),
             getPrewhereInfo(query_info_),
             data_.getPartitionValueType(),
-            virt_column_names_),
+            virt_column_names_,
+            column_names_to_return),
         .is_streaming = static_cast<bool>(create_streaming_source_)})
     , reader_settings(getMergeTreeReaderSettings(context_, query_info_))
     , prepared_parts(std::move(parts_))

src/Processors/QueryPlan/ReadFromMergeTree.h

@@ -100,7 +100,8 @@ class ReadFromMergeTree final : public ISourceStep
         Poco::Logger * log_,
         MergeTreeDataSelectAnalysisResultPtr analyzed_result_ptr_,
         bool enable_parallel_reading,
-        std::function<std::shared_ptr<ISource>(Int64 &)> create_streaming_source_ = {});
+        std::function<std::shared_ptr<ISource>(Int64 &)> create_streaming_source_ = {},
+        const Names & column_names_to_return = {});
 
     String getName() const override
     {

src/Storages/MergeTree/MergeTreeBaseSelectProcessor.cpp

@@ -633,7 +633,7 @@ void MergeTreeBaseSelectProcessor::injectVirtualColumns(
 }
 
 Block MergeTreeBaseSelectProcessor::transformHeader(
-    Block block, const PrewhereInfoPtr & prewhere_info, const DataTypePtr & partition_value_type, const Names & virtual_columns)
+    Block block, const PrewhereInfoPtr & prewhere_info, const DataTypePtr & partition_value_type, const Names & virtual_columns, const Names & columns_to_return)
 {
     if (prewhere_info)
     {
@@ -676,6 +676,11 @@ Block MergeTreeBaseSelectProcessor::transformHeader(
     }
 
     injectVirtualColumns(block, nullptr, partition_value_type, virtual_columns);
+
+    /// proton: starts. Only for read concat, we requires return columns in order
+    if (!columns_to_return.empty())
+        block.reorderColumnsInplace(columns_to_return);
+    /// proton: ends.
     return block;
 }
 

src/Storages/MergeTree/MergeTreeBaseSelectProcessor.h

@@ -50,7 +50,7 @@ class MergeTreeBaseSelectProcessor : public ISource
     ~MergeTreeBaseSelectProcessor() override;
 
     static Block transformHeader(
-        Block block, const PrewhereInfoPtr & prewhere_info, const DataTypePtr & partition_value_type, const Names & virtual_columns);
+        Block block, const PrewhereInfoPtr & prewhere_info, const DataTypePtr & partition_value_type, const Names & virtual_columns, const Names & columns_to_return = {});
 
     static std::unique_ptr<MergeTreeBlockSizePredictor> getSizePredictor(
         const MergeTreeData::DataPartPtr & data_part,

src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp

@@ -1289,6 +1289,13 @@ QueryPlanPtr MergeTreeDataSelectExecutor::readFromParts(
 
     selectColumnNames(column_names_to_return, data, real_column_names, virt_column_names, sample_factor_column_queried);
 
+    /// proton: starts. Add for read concat
+    /// Sometimes, historical source is not compatable with streaming source header
+    /// its format is always `<real columns> + <virtual columns>`, for example:
+    /// `select _tp_shard, i from t1`, historical header is `i, _tp_shard`, but streaming header is `_tp_shard, i`
+    bool required_return_columns_in_order = static_cast<bool>(create_streaming_source);
+    /// proton: ends.
+
     auto read_from_merge_tree = std::make_unique<ReadFromMergeTree>(
         std::move(parts),
         real_column_names,
@@ -1304,7 +1311,8 @@ QueryPlanPtr MergeTreeDataSelectExecutor::readFromParts(
         log,
         merge_tree_select_result_ptr,
         enable_parallel_reading,
-        std::move(create_streaming_source)
+        std::move(create_streaming_source),
+        required_return_columns_in_order ? column_names_to_return : Names{}
     );
 
     QueryPlanPtr plan = std::make_unique<QueryPlan>();

tests/stream/test_stream_smoke/0099_fixed_issues.json

@@ -572,6 +572,31 @@
             ["True"]
           ]}
         ]
-      }
+      },
+      {
+          "id": 20,
+          "tags": ["backfill", "virtual_column"],
+          "name": "#370",
+          "description": "Covering backfill from historical with virtual columns test.",
+          "steps":[
+            {
+              "statements": [
+                {"client":"python", "query_type": "table", "query": "drop stream if exists fixed_test_stream_99"},
+                {"client":"python", "query_type": "table", "exist": "fixed_test_stream_99", "exist_wait":2, "wait":1, "query": "create stream if not exists fixed_test_stream_99(id int, value int) settings shards=3, sharding_expr='weak_hash32(id)';"},
+                {"client":"python", "query_type": "table", "depends_on_stream": "fixed_test_stream_99", "wait":2, "query": "insert into fixed_test_stream_99(id, value) values(1, 2);"},
+                {"client":"python", "query_type": "stream", "depends_on_stream":"fixed_test_stream_99", "query_id":"9920", "query":"select id, _tp_shard from fixed_test_stream_99 where _tp_time > earliest_ts() settings enable_backfill_from_historical_store=true;"},
+                {"client":"python", "query_type": "table", "kill":9920, "kill_wait":3, "query": "insert into fixed_test_stream_99(id, value) values(2, 2)(3, 3);"},
+                {"client":"python", "query_type": "table", "kill":9920, "kill_wait":3, "query": "insert into fixed_test_stream_99(id, value) values(2, 2)(3, 3);"}
+              ]
+            }
+          ],
+          "expected_results": [
+            {"query_id":"9920", "expected_results":[
+              [1, 0],
+              [2, 0],
+              [3, 2]
+            ]}
+          ]
+        }
   ]
 }