Implement merge reads option (#9206)

Author: ssmike

ydb/core/kqp/common/control.cpp

@@ -0,0 +1,31 @@
+#include "control.h"
+
+#include <ydb/core/base/appdata_fwd.h>
+#include <ydb/core/control/immediate_control_board_impl.h>
+
+namespace {
+
+struct TControls {
+    std::shared_ptr<NKikimr::TControlWrapper> MergeReads;
+
+    TControls() {
+        if (auto *appData = NKikimr::AppData()) {
+            if (appData->Icb) {
+                MergeReads = std::make_shared<NKikimr::TControlWrapper>(0, 0, 1);
+                appData->Icb->RegisterSharedControl(*MergeReads,
+                    "TableServiceControls.EnableMergeDatashardReads");
+            }
+        }
+
+    }
+};
+
+}
+
+namespace NKikimr::NKqp {
+
+std::shared_ptr<TControlWrapper> MergeDatashardReadsControl() {
+    return Singleton<TControls>()->MergeReads;
+}
+
+} // namespace NKikimr::NKqp

ydb/core/kqp/common/control.h

@@ -0,0 +1,7 @@
+#include <ydb/core/control/immediate_control_board_wrapper.h>
+
+namespace NKikimr::NKqp {
+
+std::shared_ptr<TControlWrapper> MergeDatashardReadsControl();
+
+} // namespace NKikimr::NKqp

ydb/core/kqp/common/ya.make

@@ -1,6 +1,7 @@
 LIBRARY()
 
 SRCS(
+    control.cpp
     kqp_event_ids.h
     kqp_event_impl.cpp
     kqp_lwtrace_probes.cpp

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -28,6 +28,7 @@
 #include <ydb/core/kqp/opt/kqp_query_plan.h>
 #include <ydb/core/kqp/rm_service/kqp_rm_service.h>
 #include <ydb/core/grpc_services/local_rate_limiter.h>
+#include <ydb/core/kqp/common/control.h>
 
 #include <ydb/services/metadata/secret/fetcher.h>
 #include <ydb/services/metadata/secret/snapshot.h>
@@ -146,6 +147,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         , StreamResult(streamResult)
         , StatementResultIndex(statementResultIndex)
     {
+        EnableReadsMerge = *MergeDatashardReadsControl() == 1;
         TasksGraph.GetMeta().Snapshot = IKqpGateway::TKqpSnapshot(Request.Snapshot.Step, Request.Snapshot.TxId);
         TasksGraph.GetMeta().Arena = MakeIntrusive<NActors::TProtoArenaHolder>();
         TasksGraph.GetMeta().Database = Database;
@@ -1084,7 +1086,11 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         return result;
     }
 
-    TMaybe<size_t> BuildScanTasksFromSource(TStageInfo& stageInfo, const bool shardsResolved, const bool limitTasksPerNode) {
+    TMaybe<size_t> BuildScanTasksFromSource(TStageInfo& stageInfo, const bool shardsResolved, bool limitTasksPerNode) {
+        if (EnableReadsMerge) {
+            limitTasksPerNode = true;
+        }
+
         THashMap<ui64, std::vector<ui64>> nodeTasks;
         THashMap<ui64, ui64> assignedShardsCount;
 
@@ -2026,6 +2032,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
 
     ui32 StatementResultIndex;
     bool AlreadyReplied = false;
+    bool EnableReadsMerge = false;
 
 private:
     static constexpr TDuration ResourceUsageUpdateInterval = TDuration::MilliSeconds(100);

ydb/core/protos/config.proto

@@ -1475,6 +1475,14 @@ message TImmediateControlsConfig {
             DefaultValue: 0 }];
     }
 
+    message TTableServiceControls {
+        optional uint64 EnableMergeDatashardReads = 1 [(ControlOptions) = {
+            Description: "Merge reading tasks on the same node",
+            MinValue: 0,
+            MaxValue: 1,
+            DefaultValue: 0 }];
+    }
+
     optional TDataShardControls DataShardControls = 1;
     optional TTxLimitControls TxLimitControls = 2;
     optional TCoordinatorControls CoordinatorControls = 3;
@@ -1486,6 +1494,7 @@ message TImmediateControlsConfig {
     optional TDSProxyControls DSProxyControls = 9;
     optional TPDiskControls PDiskControls = 10;
     optional TBlobStorageControllerControls BlobStorageControllerControls = 11;
+    optional TTableServiceControls TableServiceControls = 12;
 };
 
 message TMeteringConfig {