Merge 23d4559 into d6fd00d

Author: gridnevvvit

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -1847,7 +1847,6 @@ class TKqpDataExecuter: public TKqpExecuterBase<TKqpDataExecuter, EExecType::Dat
                     switch (stage.GetSources(0).GetTypeCase()) {
                         case NKqpProto::TKqpSource::kReadRangesSource:
                             if (auto partitionsCount = BuildScanTasksFromSource(stageInfo,
-                                    /* shardsResolved */ StreamResult,
                                     /* limitTasksPerNode */ StreamResult)) {
                                 sourceScanPartitionsCount += *partitionsCount;
                             } else {
@@ -1914,9 +1913,6 @@ class TKqpDataExecuter: public TKqpExecuterBase<TKqpDataExecuter, EExecType::Dat
             }
         }
 
-        // For generic query all shards are already resolved
-        YQL_ENSURE(!StreamResult || remoteComputeTasks.empty());
-
         for (const auto& channel : TasksGraph.GetChannels()) {
             if (IsCrossShardChannel(TasksGraph, channel)) {
                 HasPersistentChannels = true;
@@ -1996,7 +1992,6 @@ class TKqpDataExecuter: public TKqpExecuterBase<TKqpDataExecuter, EExecType::Dat
         TasksGraph.GetMeta().UseFollowers = GetUseFollowers();
 
         if (RemoteComputeTasks) {
-            YQL_ENSURE(!StreamResult);
             TSet<ui64> shardIds;
             for (const auto& [shardId, _] : RemoteComputeTasks) {
                 shardIds.insert(shardId);
@@ -2057,6 +2052,11 @@ class TKqpDataExecuter: public TKqpExecuterBase<TKqpDataExecuter, EExecType::Dat
                 }
             }
 
+            if (shardIds.size() <= 1 && HasDatashardSourceScan) {
+                // nothing to merge
+                HasDatashardSourceScan = false;
+            }
+
             if ((HasOlapTable || HasDatashardSourceScan) && shardIds) {
                 LOG_D("Start resolving tablets nodes... (" << shardIds.size() << ")");
                 ExecuterStateSpan = NWilson::TSpan(

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -213,6 +213,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         auto& reply = *ev->Get();
 
         KqpShardsResolverId = {};
+        ShardsResolved = true;
 
         // TODO: count resolve time in CpuTime
 
@@ -419,7 +420,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
                     auto& extraData = ExtraData[computeActor];
                     extraData.TaskId = taskId;
                     extraData.Data.Swap(state.MutableExtraData());
-                    
+
 
                     Stats->AddComputeActorStats(
                         computeActor.NodeId(),
@@ -1015,7 +1016,7 @@ 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 limitTasksPerNode) {
         THashMap<ui64, std::vector<ui64>> nodeTasks;
         THashMap<ui64, ui64> assignedShardsCount;
 
@@ -1050,7 +1051,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
             if (nodeId) {
                 task.Meta.NodeId = *nodeId;
             } else {
-                YQL_ENSURE(!shardsResolved);
+                YQL_ENSURE(!ShardsResolved);
                 task.Meta.ShardId = taskLocation;
             }
             FillSecureParamsFromStage(task.Meta.SecureParams, stage);
@@ -1144,15 +1145,14 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
                 ? TMaybe<ui64>{*nodeIdPtr}
                 : Nothing();
 
-            YQL_ENSURE(!shardsResolved || nodeId);
+            YQL_ENSURE(!ShardsResolved || nodeId);
             YQL_ENSURE(Stats);
 
             if (shardId) {
                 Stats->AffectedShards.insert(*shardId);
             }
 
-            if (limitTasksPerNode) {
-                YQL_ENSURE(shardsResolved);
+            if (limitTasksPerNode && ShardsResolved) {
                 const auto maxScanTasksPerNode = GetScanTasksPerNode(stageInfo, /* isOlapScan */ false, *nodeId);
                 auto& nodeTasks = nodeIdToTasks[*nodeId];
                 if (nodeTasks.size() < maxScanTasksPerNode) {
@@ -1925,6 +1925,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
     TActorId Target;
     ui64 TxId = 0;
 
+    bool ShardsResolved = false;
     TKqpTasksGraph TasksGraph;
 
     TActorId KqpTableResolverId;

ydb/core/kqp/executer_actor/kqp_scan_executer.cpp

@@ -184,7 +184,6 @@ class TKqpScanExecuter : public TKqpExecuterBase<TKqpScanExecuter, EExecType::Sc
                     case NKqpProto::TKqpSource::kReadRangesSource:
                         BuildScanTasksFromSource(
                             stageInfo,
-                            /* shardsResolved */ true,
                             /* limitTasksPerNode */ false);
                         break;
                     default:

ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp

@@ -2,6 +2,7 @@
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 #include <ydb/core/kqp/ut/common/columnshard.h>
 #include <ydb/core/testlib/common_helper.h>
+#include <ydb/core/base/tablet_pipecache.h>
 #include <ydb/core/tx/columnshard/hooks/abstract/abstract.h>
 #include <ydb/core/tx/columnshard/hooks/testing/controller.h>
 #include <ydb/public/lib/ut_helpers/ut_helpers_query.h>
@@ -237,6 +238,46 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
         }
     }
 
+    Y_UNIT_TEST(Followers) {
+        NKikimrConfig::TAppConfig config;
+
+        auto kikimr = TKikimrRunner(TKikimrSettings().SetAppConfig(config).SetUseRealThreads(false));
+        auto db = kikimr.GetQueryClient();
+
+        TExecuteQuerySettings settings;
+
+        THashMap<TActorId, ui64> countResolveTablet;
+        countResolveTablet[NKikimr::MakePipePerNodeCacheID(false)] = 0;
+        countResolveTablet[NKikimr::MakePipePerNodeCacheID(true)] = 0;
+        auto counterObserver = [&countResolveTablet](TAutoPtr<IEventHandle>& ev) -> auto {
+            if (ev->GetTypeRewrite() == NKikimr::TEvPipeCache::TEvGetTabletNode::EventType) {
+                countResolveTablet[ev->Recipient]++;
+                return TTestActorRuntime::EEventAction::PROCESS;
+            }
+
+            return TTestActorRuntime::EEventAction::PROCESS;
+        };
+
+        kikimr.GetTestServer().GetRuntime()->SetObserverFunc(counterObserver);
+
+        Y_DEFER {
+            kikimr.GetTestServer().GetRuntime()->SetObserverFunc(TTestActorRuntime::DefaultObserverFunc);
+        };
+
+        {
+            const TString query = "SELECT Key, Value2 FROM TwoShard WHERE Key = 1u";
+            auto result = kikimr.RunCall([&] { return db.ExecuteQuery(query, NQuery::TTxControl::BeginTx(TTxSettings::StaleRO()).CommitTx()).ExtractValueSync(); });
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
+            CompareYson(R"([
+                [[1u];[-1]]
+            ])", FormatResultSetYson(result.GetResultSet(0)));
+            UNIT_ASSERT(countResolveTablet[NKikimr::MakePipePerNodeCacheID(false)] == 0);
+            // using followers resolver.
+            UNIT_ASSERT(countResolveTablet[NKikimr::MakePipePerNodeCacheID(true)] > 0);
+        }
+    }
+
     Y_UNIT_TEST(ExecuteQueryWithWorkloadManager) {
         NKikimrConfig::TAppConfig config;
         config.MutableFeatureFlags()->SetEnableResourcePools(true);