Merge 88f7ac2 into 87530be

Author: nikvas0

ydb/core/kqp/common/compilation/events.h

@@ -152,14 +152,25 @@ struct TEvParseResponse: public TEventLocal<TEvParseResponse, TKqpEvents::EvPars
 };
 
 struct TEvSplitResponse: public TEventLocal<TEvSplitResponse, TKqpEvents::EvSplitResponse> {
-    TEvSplitResponse(const TKqpQueryId& query, TVector<NYql::TExprNode::TPtr> exprs, NYql::TExprNode::TPtr world, THolder<NYql::TExprContext> ctx)
-        : Query(query)
+    TEvSplitResponse(
+            Ydb::StatusIds::StatusCode status,
+            const NYql::TIssues& issues,
+            const TKqpQueryId& query,
+            TVector<NYql::TExprNode::TPtr> exprs,
+            NYql::TExprNode::TPtr world,
+            std::shared_ptr<NYql::TExprContext> ctx)
+        : Status(status)
+        , Issues(issues)
+        , Query(query)
         , Ctx(std::move(ctx))
         , Exprs(std::move(exprs))
         , World(std::move(world)) {}
 
+    Ydb::StatusIds::StatusCode Status;
+    NYql::TIssues Issues;
+
     TKqpQueryId Query;
-    THolder<NYql::TExprContext> Ctx;
+    std::shared_ptr<NYql::TExprContext> Ctx;
     TVector<NYql::TExprNode::TPtr> Exprs;
     NYql::TExprNode::TPtr World;
 };

ydb/core/kqp/compile_service/kqp_compile_actor.cpp

@@ -127,7 +127,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
     STFUNC(CompileState) {
         try {
             switch (ev->GetTypeRewrite()) {
-                HFunc(TEvKqp::TEvContinueProcess, Handle);
+                HFunc(TEvKqp::TEvContinueProcess, HandleCompile);
                 cFunc(TEvents::TSystem::Wakeup, HandleTimeout);
             default:
                 UnexpectedEvent("CompileState", ev->GetTypeRewrite());
@@ -137,6 +137,20 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
         }
     }
 
+    STFUNC(SplitState) {
+        try {
+            switch (ev->GetTypeRewrite()) {
+                HFunc(TEvKqp::TEvContinueProcess, HandleSplit);
+                cFunc(TEvents::TSystem::Wakeup, HandleTimeout);
+            default:
+                UnexpectedEvent("SplitState", ev->GetTypeRewrite());
+            }
+        } catch (const yexception& e) {
+            InternalError(e.what());
+        }
+    }
+
+
 private:
     TVector<TQueryAst> GetAstStatements(const TActorContext &ctx) {
         TString cluster = QueryId.Cluster;
@@ -158,9 +172,11 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
         ALOG_DEBUG(NKikimrServices::KQP_COMPILE_ACTOR, "Send split result"
             << ", self: " << SelfId()
             << ", owner: " << Owner
-            << (!result.Exprs.empty() ? ", split is successful" : ", split is not successful"));
+            << ", success: " << GetYdbStatus(result)
+            << ", issues: " << result.Issues().ToOneLineString());
 
         auto responseEv = MakeHolder<TEvKqp::TEvSplitResponse>(
+            GetYdbStatus(result), result.Issues(),
             QueryId, std::move(result.Exprs), std::move(result.World), std::move(result.Ctx));
         Send(Owner, responseEv.Release());
 
@@ -174,11 +190,38 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
     }
 
     void StartSplitting(const TActorContext &ctx) {
+        Become(&TKqpCompileActor::SplitState);
+        TimeoutTimerActorId = CreateLongTimer(ctx, CompilationTimeout, new IEventHandle(SelfId(), SelfId(),
+            new TEvents::TEvWakeup()));
+
         const auto prepareSettings = PrepareCompilationSettings(ctx);
-        auto result = KqpHost->SplitQuery(QueryRef, prepareSettings);
+        AsyncSplitResult = KqpHost->SplitQuery(QueryRef, prepareSettings);
+        ContinueSplittig(ctx);
+    }
 
-        Become(&TKqpCompileActor::CompileState);
-        ReplySplitResult(ctx, std::move(result));
+    void ContinueSplittig(const TActorContext &ctx) {
+        TActorSystem* actorSystem = ctx.ExecutorThread.ActorSystem;
+        TActorId selfId = ctx.SelfID;
+
+        auto callback = [actorSystem, selfId](const TFuture<bool>& future) {
+            bool finished = future.GetValue();
+            auto processEv = MakeHolder<TEvKqp::TEvContinueProcess>(0, finished);
+            actorSystem->Send(selfId, processEv.Release());
+        };
+
+        AsyncSplitResult->Continue().Apply(callback);
+    }
+
+    void HandleSplit(TEvKqp::TEvContinueProcess::TPtr &ev, const TActorContext &ctx) {
+        Y_ENSURE(!ev->Get()->QueryId);
+
+        if (!ev->Get()->Finished) {
+            ContinueSplittig(ctx);
+            return;
+        }
+
+        auto splitResult = std::move(AsyncSplitResult->GetResult());
+        ReplySplitResult(ctx, std::move(splitResult));
     }
 
     void StartParsing(const TActorContext &ctx) {
@@ -460,7 +503,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
         KqpCompileResult->AllowCache = CanCacheQuery(KqpCompileResult->PreparedQuery->GetPhysicalQuery()) && allowCache;
     }
 
-    void Handle(TEvKqp::TEvContinueProcess::TPtr &ev, const TActorContext &ctx) {
+    void HandleCompile(TEvKqp::TEvContinueProcess::TPtr &ev, const TActorContext &ctx) {
         Y_ENSURE(!ev->Get()->QueryId);
 
         TYqlLogScope logScope(ctx, NKikimrServices::KQP_YQL, YqlName, UserRequestContext->TraceId);
@@ -558,6 +601,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
     TIntrusivePtr<IKqpGateway> Gateway;
     TIntrusivePtr<IKqpHost> KqpHost;
     TIntrusivePtr<IKqpHost::IAsyncQueryResult> AsyncCompileResult;
+    TIntrusivePtr<IKqpHost::IAsyncSplitResult> AsyncSplitResult;
     std::shared_ptr<TKqpCompileResult> KqpCompileResult;
     std::optional<TString> ReplayMessage;  // here metadata is encoded protobuf - for logs
     std::optional<TString> ReplayMessageUserView;  // here metadata is part of json - full readable json for diagnostics