enable fan out for the same source

Author: yl-lisen

src/Storages/Streaming/SourceColumnsDescription.cpp

@@ -30,21 +30,31 @@ void SourceColumnsDescription::PhysicalColumnPositions::clear()
     subcolumns.clear();
 }
 
-SourceColumnsDescription::SourceColumnsDescription(const Names & required_column_names, StorageSnapshotPtr storage_snapshot)
+SourceColumnsDescription::SourceColumnsDescription(
+    const Names & required_column_names, StorageSnapshotPtr storage_snapshot, bool enable_partial_read)
     : SourceColumnsDescription(
-        storage_snapshot->getColumnsByNames(GetColumnsOptions(GetColumnsOptions::All).withSubcolumns().withVirtuals().withExtendedObjects(), required_column_names),
+        storage_snapshot->getColumnsByNames(
+            GetColumnsOptions(GetColumnsOptions::All).withSubcolumns().withVirtuals().withExtendedObjects(), required_column_names),
         storage_snapshot->getMetadataForQuery()->getSampleBlock(),
-        storage_snapshot->getColumns(GetColumnsOptions(GetColumnsOptions::All).withExtendedObjects()))
+        storage_snapshot->getColumns(GetColumnsOptions(GetColumnsOptions::All).withExtendedObjects()),
+        enable_partial_read)
 {
 }
 
-SourceColumnsDescription::SourceColumnsDescription(const NamesAndTypesList & columns_to_read, const Block & schema, const NamesAndTypesList & all_extended_columns)
+SourceColumnsDescription::SourceColumnsDescription(
+    const NamesAndTypesList & columns_to_read,
+    const Block & schema,
+    const NamesAndTypesList & all_extended_columns,
+    bool enable_partial_read)
 {
     /// FIXME, when we have multi-version of schema, the header and the schema may be mismatched
     auto column_size = columns_to_read.size();
 
+    /// Just read required partial physical columns
+    if (enable_partial_read)
+        physical_column_positions_to_read.positions.reserve(column_size);
+
     positions.reserve(column_size);
-    physical_column_positions_to_read.positions.reserve(column_size);
     subcolumns_to_read.reserve(column_size);
 
     std::vector<uint16_t> read_all_subcolumns_positions;
@@ -112,45 +122,48 @@ SourceColumnsDescription::SourceColumnsDescription(const NamesAndTypesList & col
             auto pos_in_schema = schema.getPositionByName(name_in_storage);
             const auto & column_in_storage = schema.getByName(name_in_storage);
 
-            /// Calculate main column pos
-            size_t physical_pos_in_schema_to_read = 0;
-            /// We don't need to read duplicate physical columns from schema
-            auto physical_pos_iter = std::find(
-                physical_column_positions_to_read.positions.begin(), physical_column_positions_to_read.positions.end(), pos_in_schema);
-            if (physical_pos_iter == physical_column_positions_to_read.positions.end())
+            size_t physical_pos_in_schema_to_read = pos_in_schema;
+            /// Specially, re-calculate pos in partially read schema
+            if (enable_partial_read)
             {
-                physical_pos_in_schema_to_read = physical_column_positions_to_read.positions.size();
-                physical_column_positions_to_read.positions.emplace_back(pos_in_schema);
-
-                /// json, array(json), tuple(..., json, ...)
-                if (column_in_storage.type->hasDynamicSubcolumns())
+                /// We don't need to read duplicate physical columns from schema
+                auto physical_pos_iter = std::find(
+                    physical_column_positions_to_read.positions.begin(), physical_column_positions_to_read.positions.end(), pos_in_schema);
+                if (physical_pos_iter == physical_column_positions_to_read.positions.end())
                 {
-                    /// We like to read parent json column once if multiple subcolumns of the same json are required
-                    /// like `select json.a, json.b from stream`
-                    auto find_iter = std::find_if(
-                        physical_object_columns_to_read.begin(),
-                        physical_object_columns_to_read.end(),
-                        [&column](const auto & col_name_type) { return col_name_type.name == column.name; });
-
-                    if (find_iter == physical_object_columns_to_read.end())
+                    physical_pos_in_schema_to_read = physical_column_positions_to_read.positions.size();
+                    physical_column_positions_to_read.positions.emplace_back(pos_in_schema);
+
+                    /// json, array(json), tuple(..., json, ...)
+                    if (column_in_storage.type->hasDynamicSubcolumns())
                     {
-                        if (column.isSubcolumn())
+                        /// We like to read parent json column once if multiple subcolumns of the same json are required
+                        /// like `select json.a, json.b from stream`
+                        auto find_iter = std::find_if(
+                            physical_object_columns_to_read.begin(),
+                            physical_object_columns_to_read.end(),
+                            [&column](const auto & col_name_type) { return col_name_type.name == column.name; });
+
+                        if (find_iter == physical_object_columns_to_read.end())
                         {
-                            /// When reading a subcolumn of a json like `select json.a from stream`, we will need read the parent `json` column
-                            auto name_and_type = all_extended_columns.tryGetByName(name_in_storage);
-                            assert(name_and_type);
-                            physical_object_columns_to_read.emplace_back(std::move(*name_and_type));
-                        }
-                        else
-                        {
-                            /// This column is parent json column, like `select json from stream`, use the name and type directly
-                            physical_object_columns_to_read.emplace_back(column);
+                            if (column.isSubcolumn())
+                            {
+                                /// When reading a subcolumn of a json like `select json.a from stream`, we will need read the parent `json` column
+                                auto name_and_type = all_extended_columns.tryGetByName(name_in_storage);
+                                assert(name_and_type);
+                                physical_object_columns_to_read.emplace_back(std::move(*name_and_type));
+                            }
+                            else
+                            {
+                                /// This column is parent json column, like `select json from stream`, use the name and type directly
+                                physical_object_columns_to_read.emplace_back(column);
+                            }
                         }
                     }
                 }
+                else
+                    physical_pos_in_schema_to_read = physical_pos_iter - physical_column_positions_to_read.positions.begin();
             }
-            else
-                physical_pos_in_schema_to_read = physical_pos_iter - physical_column_positions_to_read.positions.begin();
 
             /// For subcolumn, which dependents on the main column
             if (column.isSubcolumn())
@@ -181,7 +194,7 @@ SourceColumnsDescription::SourceColumnsDescription(const NamesAndTypesList & col
         physical_column_positions_to_read.subcolumns.erase(pos);
 
     /// Clients like to read virtual columns only, add `_tp_time`, then we know how many rows
-    if (physical_column_positions_to_read.positions.empty())
+    if (enable_partial_read && physical_column_positions_to_read.positions.empty())
         physical_column_positions_to_read.positions.emplace_back(schema.getPositionByName(ProtonConsts::RESERVED_EVENT_TIME));
 }
 }

src/Storages/Streaming/SourceColumnsDescription.h

@@ -21,8 +21,12 @@ using StorageSnapshotPtr = std::shared_ptr<StorageSnapshot>;
 struct SourceColumnsDescription
 {
     SourceColumnsDescription() = default;
-    SourceColumnsDescription(const Names & required_column_names, StorageSnapshotPtr storage_snapshot);
-    SourceColumnsDescription(const NamesAndTypesList & columns_to_read, const Block & schema, const NamesAndTypesList & all_extended_columns);
+    SourceColumnsDescription(const Names & required_column_names, StorageSnapshotPtr storage_snapshot, bool enable_partial_read = true);
+    SourceColumnsDescription(
+        const NamesAndTypesList & columns_to_read,
+        const Block & schema,
+        const NamesAndTypesList & all_extended_columns,
+        bool enable_partial_read = true);
 
     enum class ReadColumnType : uint8_t
     {

src/Storages/Streaming/StorageStream.cpp

@@ -601,22 +601,18 @@ void StorageStream::readStreaming(
 
     assert(query_info.seek_to_info);
     const auto & settings_ref = context_->getSettingsRef();
-    /// 1) Checkpointed queries shall not be multiplexed
-    /// 2) Queries which seek to a specific timestamp shall not be multiplexed
-    auto share_resource_group = (settings_ref.query_resource_group.value == "shared")
-        && (query_info.seek_to_info->getSeekTo().empty() || query_info.seek_to_info->getSeekTo() == "latest")
-        && (settings_ref.exec_mode == ExecuteMode::NORMAL);
-
-    if (share_resource_group)
+    if (settings_ref.query_resource_group.value == "shared")
     {
+        auto offsets = stream_shards.back()->getOffsets(query_info.seek_to_info);
         for (auto stream_shard : shards_to_read)
         {
+            const auto & offset = offsets[stream_shard->shard];
             if (!column_names.empty())
                 pipes.emplace_back(
-                    stream_shard->source_multiplexers->createChannel(stream_shard->shard, column_names, storage_snapshot, context_));
+                    source_multiplexers->createChannel(std::move(stream_shard), column_names, storage_snapshot, context_, offset));
             else
-                pipes.emplace_back(stream_shard->source_multiplexers->createChannel(
-                    stream_shard->shard, {ProtonConsts::RESERVED_EVENT_TIME}, storage_snapshot, context_));
+                pipes.emplace_back(source_multiplexers->createChannel(
+                    std::move(stream_shard), {ProtonConsts::RESERVED_EVENT_TIME}, storage_snapshot, context_, offset));
         }
 
         LOG_INFO(log, "Starting reading {} streams in shared resource group", pipes.size());
@@ -945,6 +941,8 @@ void StorageStream::startup()
         assert(native_log->enabled());
     }
 
+    source_multiplexers.reset(new StreamingStoreSourceMultiplexers(getContext(), log));
+
     log_initialized.test_and_set();
 
     LOG_INFO(log, "Started");

src/Storages/Streaming/StorageStream.h

@@ -10,6 +10,7 @@
 
 #include <Storages/MergeTree/MergeTreeData.h>
 #include <Storages/MergeTree/MergeTreeMutationStatus.h>
+#include <Storages/Streaming/StreamingStoreSourceMultiplexer.h>
 
 namespace nlog
 {
@@ -298,8 +299,7 @@ class StorageStream final : public shared_ptr_helper<StorageStream>, public Merg
         UInt64 base_block_id,
         UInt64 sub_block_id);
 
-    void
-    appendToNativeLog(nlog::RecordPtr & record, IngestMode /*ingest_mode*/, klog::AppendCallback callback, klog::CallbackData data);
+    void appendToNativeLog(nlog::RecordPtr & record, IngestMode /*ingest_mode*/, klog::AppendCallback callback, klog::CallbackData data);
 
     void appendToKafka(
         nlog::RecordPtr & record,
@@ -354,5 +354,8 @@ class StorageStream final : public shared_ptr_helper<StorageStream>, public Merg
 
     std::atomic_flag inited;
     std::atomic_flag stopped;
+
+    /// Multiplex latest records of each shard.
+    std::unique_ptr<StreamingStoreSourceMultiplexers> source_multiplexers;
 };
 }

src/Storages/Streaming/StreamShard.cpp

@@ -124,8 +124,6 @@ StreamShard::~StreamShard()
 
 void StreamShard::startup()
 {
-    source_multiplexers.reset(new StreamingStoreSourceMultiplexers(shared_from_this(), storage_stream->getContext(), log));
-
     initLog();
 
     /// for virtual tables or in-memory storage type, there is no storage object

src/Storages/Streaming/StreamShard.h

@@ -148,8 +148,6 @@ class StreamShard final : public std::enable_shared_from_this<StreamShard>
 
     std::unique_ptr<StreamCallbackData> callback_data;
 
-    std::unique_ptr<StreamingStoreSourceMultiplexers> source_multiplexers;
-
     // For random shard index generation
     mutable std::mutex rng_mutex;
     pcg64 rng;

src/Storages/Streaming/StreamingBlockReaderNativeLog.cpp

@@ -148,9 +148,8 @@ nlog::RecordPtrs StreamingBlockReaderNativeLog::processCached(nlog::RecordPtrs r
             {
                 /// In general, an object has a large number of subcolumns,
                 /// so when a few subcolumns required for the object, we only copy partials to improve performance
-                if (isObject(col_with_type->type) && !schema_ctx.column_positions.positions.empty())
+                if (isObject(col_with_type->type) && !schema_ctx.column_positions.subcolumns.empty())
                 {
-                    assert(column_names.size() == schema_ctx.column_positions.positions.size());
                     auto iter = schema_ctx.column_positions.subcolumns.find(i);
                     if (iter != schema_ctx.column_positions.subcolumns.end())
                     {

src/Storages/Streaming/StreamingStoreSource.cpp

@@ -1,7 +1,7 @@
-#include "StreamingStoreSource.h"
-#include "StreamShard.h"
-#include "StreamingBlockReaderKafka.h"
-#include "StreamingBlockReaderNativeLog.h"
+#include <Storages/Streaming/StreamShard.h>
+#include <Storages/Streaming/StreamingBlockReaderKafka.h>
+#include <Storages/Streaming/StreamingBlockReaderNativeLog.h>
+#include <Storages/Streaming/StreamingStoreSource.h>
 
 #include <Interpreters/inplaceBlockConversions.h>
 #include <KafkaLog/KafkaWALPool.h>
@@ -16,7 +16,8 @@ StreamingStoreSource::StreamingStoreSource(
     ContextPtr context_,
     Int64 sn,
     Poco::Logger * log_)
-    : StreamingStoreSourceBase(header, storage_snapshot_, std::move(context_), log_, ProcessorID::StreamingStoreSourceID)
+    : StreamingStoreSourceBase(
+        header, storage_snapshot_, /*enable_partial_read*/ true, std::move(context_), log_, ProcessorID::StreamingStoreSourceID)
 {
     const auto & settings = query_context->getSettingsRef();
     if (settings.record_consume_batch_count.value != 0)
@@ -98,7 +99,8 @@ void StreamingStoreSource::readAndProcess()
                     /// NOTE: The `FilterTransform` will try optimizing filter ConstColumn to always_false or always_true,
                     /// for exmaple: `_tp_sn < 1`, if filter first data _tp_sn is 0, it will be optimized always_true.
                     /// So we can not create a constant column, since the virtual column data isn't constants value in fact.
-                    auto virtual_column = columns_desc.virtual_col_types[pos.virtualPosition()]->createColumnConst(rows, ts)->convertToFullColumnIfConst();
+                    auto virtual_column
+                        = columns_desc.virtual_col_types[pos.virtualPosition()]->createColumnConst(rows, ts)->convertToFullColumnIfConst();
                     columns.push_back(std::move(virtual_column));
                     break;
                 }

src/Storages/Streaming/StreamingStoreSource.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "StreamingStoreSourceBase.h"
+#include <Storages/Streaming/StreamingStoreSourceBase.h>
 
 namespace DB
 {

src/Storages/Streaming/StreamingStoreSourceBase.cpp

@@ -1,4 +1,4 @@
-#include "StreamingStoreSourceBase.h"
+#include <Storages/Streaming/StreamingStoreSourceBase.h>
 
 #include <Checkpoint/CheckpointContext.h>
 #include <Checkpoint/CheckpointCoordinator.h>
@@ -16,14 +16,19 @@ extern const int RECOVER_CHECKPOINT_FAILED;
 }
 
 StreamingStoreSourceBase::StreamingStoreSourceBase(
-    const Block & header, const StorageSnapshotPtr & storage_snapshot_, ContextPtr query_context_, Poco::Logger * log_, ProcessorID pid_)
+    const Block & header,
+    const StorageSnapshotPtr & storage_snapshot_,
+    bool enable_partial_read,
+    ContextPtr query_context_,
+    Poco::Logger * log_,
+    ProcessorID pid_)
     : ISource(header, true, pid_)
     , storage_snapshot(
           std::make_shared<StorageSnapshot>(*storage_snapshot_)) /// We like to make a copy of it since we will mutate the snapshot
     , query_context(std::move(query_context_))
     , log(log_)
     , header_chunk(header.getColumns(), 0)
-    , columns_desc(header.getNames(), storage_snapshot)
+    , columns_desc(header.getNames(), storage_snapshot, enable_partial_read)
 {
     is_streaming = true;
 
@@ -52,7 +57,8 @@ StreamingStoreSourceBase::getSubcolumnFromBlock(const Block & block, size_t pare
 
             /// Convert subcolumn if the subcolumn type of dynamic object may be dismatched with header.
             /// FIXME: Cache the ExpressionAction
-            Block subcolumn_block({ColumnWithTypeAndName{std::move(subcolumn), std::move(subcolumn_type), subcolumn_pair.name}}); /// NOLINT(performance-move-const-arg)
+            Block subcolumn_block({ColumnWithTypeAndName{
+                std::move(subcolumn), std::move(subcolumn_type), subcolumn_pair.name}}); /// NOLINT(performance-move-const-arg)
             ExpressionActions convert_act(ActionsDAG::makeConvertingActions(
                 subcolumn_block.getColumnsWithTypeAndName(),
                 {ColumnWithTypeAndName{subcolumn_pair.type->createColumn(), subcolumn_pair.type, subcolumn_pair.name}},