Merge aea4a05 into 50b9d76

Author: GrigoriyPA

ydb/core/fq/libs/compute/ydb/result_writer_actor.cpp

@@ -1,5 +1,6 @@
 #include "base_compute_actor.h"
 
+#include <ydb/core/fq/libs/common/rows_proto_splitter.h>
 #include <ydb/core/fq/libs/common/util.h>
 #include <ydb/core/fq/libs/compute/common/metrics.h>
 #include <ydb/core/fq/libs/compute/common/retry_actor.h>
@@ -22,6 +23,8 @@
 #include <ydb/library/actors/core/log.h>
 #include <library/cpp/protobuf/interop/cast.h>
 
+#include <queue>
+
 #define LOG_E(stream) LOG_ERROR_S(*TlsActivationContext, NKikimrServices::FQ_RUN_ACTOR, "[ydb] [ResultWriter] QueryId: " << Params.QueryId << " OperationId: " << ProtoToString(OperationId) << " " << stream)
 #define LOG_W(stream) LOG_WARN_S( *TlsActivationContext, NKikimrServices::FQ_RUN_ACTOR, "[ydb] [ResultWriter] QueryId: " << Params.QueryId << " OperationId: " << ProtoToString(OperationId) << " " << stream)
 #define LOG_I(stream) LOG_INFO_S( *TlsActivationContext, NKikimrServices::FQ_RUN_ACTOR, "[ydb] [ResultWriter] QueryId: " << Params.QueryId << " OperationId: " << ProtoToString(OperationId) << " " << stream)
@@ -100,25 +103,36 @@ class TResultSetWriterActor : public TBaseComputeActor<TResultSetWriterActor> {
             return;
         }
 
-        auto startTime = TInstant::Now();
+        LOG_I("ResultSetId: " << ResultSetId << " FetchToken: " << FetchToken << " Successfully fetched " << response.ResultSet->RowsCount() << " rows");
         Truncated |= response.ResultSet->Truncated();
         FetchToken = response.NextFetchToken;
         auto emptyResultSet = response.ResultSet->RowsCount() == 0;
-        const auto resultSetProto = NYdb::TProtoAccessor::GetProto(*response.ResultSet);
+        auto resultSetProto = NYdb::TProtoAccessor::GetProto(*response.ResultSet);
 
         if (!emptyResultSet) {
-            auto chunk = CreateProtoRequestWithoutResultSet(Offset);
-            WriterInflight[Cookie] = {Offset, startTime};
-            Offset += response.ResultSet->RowsCount();
-            *chunk.mutable_result_set() = resultSetProto;
-            auto writeResultCounters = Counters.GetCounters(ERequestType::RT_WRITE_RESULT_SET);
-            writeResultCounters->InFly->Inc();
-            Send(NFq::MakeInternalServiceActorId(), new NFq::TEvInternalService::TEvWriteResultRequest(std::move(chunk)), 0, Cookie++);
+            NFq::TRowsProtoSplitter rowsSplitter(std::move(resultSetProto), ProtoMessageLimit, BaseProtoByteSize, MaxRowsCountPerChunk);
+            auto splittedResultSets = rowsSplitter.Split();
+
+            if (!splittedResultSets.Success) {
+                LOG_E("ResultSetId: " << ResultSetId << " Can't split script result: " << splittedResultSets.Issues.ToOneLineString());
+                Send(Parent, new TEvYdbCompute::TEvResultSetWriterResponse(ResultSetId, splittedResultSets.Issues, NYdb::EStatus::INTERNAL_ERROR));
+                FailedAndPassAway();
+                return;
+            }
+
+            for (auto& resultSet : splittedResultSets.ResultSets) {
+                auto protoReq = CreateProtoRequestWithoutResultSet(Offset);
+                Offset += resultSet.rows().size();
+                protoReq.mutable_result_set()->Swap(&resultSet);
+                ResultChunks.emplace(std::move(protoReq));
+            }
+
+            TryStartResultWriters();
         }
 
         if (WriterInflight.empty()) {
             SendReplyAndPassAway();
-        } else if (FetchToken && WriterInflight.size() < MAX_WRITER_INFLIGHT) {
+        } else if (FetchToken && (WriterInflight.size() + ResultChunks.size()) < 2 * MAX_WRITER_INFLIGHT) {
             SendFetchScriptResultRequest();
         }
     }
@@ -149,17 +163,31 @@ class TResultSetWriterActor : public TBaseComputeActor<TResultSetWriterActor> {
             return;
         }
 
+        TryStartResultWriters();
+
         writeResultCounters->Ok->Inc();
         LOG_I("ResultSetId: " << ResultSetId << " Cookie: " << cookie << " Result successfully written for offset " << meta.Offset);
         if (FetchToken) {
-            if (FetchToken != LastProcessedToken) {
+            if (FetchToken != LastProcessedToken && (WriterInflight.size() + ResultChunks.size()) < 2 * MAX_WRITER_INFLIGHT) {
                 SendFetchScriptResultRequest();
             }
         } else if (WriterInflight.empty()) {
             SendReplyAndPassAway();
         }
     }
 
+    void TryStartResultWriters() {
+        auto writeResultCounters = Counters.GetCounters(ERequestType::RT_WRITE_RESULT_SET);
+        while (!ResultChunks.empty() && WriterInflight.size() < MAX_WRITER_INFLIGHT) {
+            auto chunk = std::move(ResultChunks.front());
+            ResultChunks.pop();
+
+            WriterInflight[Cookie] = {static_cast<int64_t>(chunk.offset()), TInstant::Now()};
+            writeResultCounters->InFly->Inc();
+            Send(NFq::MakeInternalServiceActorId(), new NFq::TEvInternalService::TEvWriteResultRequest(std::move(chunk)), 0, Cookie++);
+        }
+    }
+
     void SendFetchScriptResultRequest() {
         LastProcessedToken = FetchToken;
         Register(new TRetryActor<TEvYdbCompute::TEvFetchScriptResultRequest, TEvYdbCompute::TEvFetchScriptResultResponse, NKikimr::NOperationId::TOperationId, int64_t, TString>(Counters.GetCounters(ERequestType::RT_FETCH_SCRIPT_RESULT), SelfId(), Connector, OperationId, ResultSetId, FetchToken));
@@ -194,6 +222,10 @@ class TResultSetWriterActor : public TBaseComputeActor<TResultSetWriterActor> {
     bool Truncated = false;
     TString FetchToken;
     TString LastProcessedToken;
+    const size_t ProtoMessageLimit = 10_MB;
+    const size_t MaxRowsCountPerChunk = 100'000;
+    const size_t BaseProtoByteSize = CreateProtoRequestWithoutResultSet(0).ByteSizeLong();
+    std::queue<Fq::Private::WriteTaskResultRequest> ResultChunks;
 };
 
 class TResultWriterActor : public TBaseComputeActor<TResultWriterActor> {

ydb/core/grpc_services/query/rpc_fetch_script_results.cpp

@@ -25,6 +25,7 @@ using TEvFetchScriptResultsRequest = TGrpcRequestNoOperationCall<Ydb::Query::Fet
     Ydb::Query::FetchScriptResultsResponse>;
 
 constexpr i64 MAX_ROWS_LIMIT = 1000;
+constexpr i64 MAX_SIZE_LIMIT = 60_MB;
 
 class TFetchScriptResultsRPC : public TRpcRequestActor<TFetchScriptResultsRPC, TEvFetchScriptResultsRequest, false> {
 public:
@@ -45,7 +46,7 @@ class TFetchScriptResultsRPC : public TRpcRequestActor<TFetchScriptResultsRPC, T
             return;
         }
 
-        if (req->rows_limit() <= 0) {
+        if (req->rows_limit() < 0) {
             Reply(Ydb::StatusIds::BAD_REQUEST, "Invalid rows limit");
             return;
         }
@@ -70,28 +71,27 @@ class TFetchScriptResultsRPC : public TRpcRequestActor<TFetchScriptResultsRPC, T
             return;
         }
 
-        Register(NKqp::CreateGetScriptExecutionResultActor(SelfId(), DatabaseName, ExecutionId, req->result_set_index(), RowsOffset, req->rows_limit() + 1));
+        Register(NKqp::CreateGetScriptExecutionResultActor(SelfId(), DatabaseName, ExecutionId, req->result_set_index(), RowsOffset, req->rows_limit(), req->rows_limit() ? 0 : MAX_SIZE_LIMIT, Request->GetDeadline()));
 
         Become(&TFetchScriptResultsRPC::StateFunc);
     }
 
 private:
     STRICT_STFUNC(StateFunc,
-        hFunc(NKqp::TEvKqp::TEvFetchScriptResultsResponse, Handle);
+        hFunc(NKqp::TEvFetchScriptResultsResponse, Handle);
     )
 
-    void Handle(NKqp::TEvKqp::TEvFetchScriptResultsResponse::TPtr& ev) {
+    void Handle(NKqp::TEvFetchScriptResultsResponse::TPtr& ev) {
         Ydb::Query::FetchScriptResultsResponse resp;
-        resp.set_status(ev->Get()->Record.GetStatus());
-        resp.mutable_issues()->Swap(ev->Get()->Record.MutableIssues());
-        resp.set_result_set_index(static_cast<i64>(ev->Get()->Record.GetResultSetIndex()));
-        if (ev->Get()->Record.HasResultSet()) {
-            resp.mutable_result_set()->Swap(ev->Get()->Record.MutableResultSet());
-
-            const auto* userReq = GetProtoRequest();
-            if (resp.mutable_result_set()->rows_size() == userReq->rows_limit() + 1) {
-                resp.mutable_result_set()->mutable_rows()->DeleteSubrange(userReq->rows_limit(), 1);
-                resp.set_next_fetch_token(ToString(RowsOffset + userReq->rows_limit()));
+        resp.set_status(ev->Get()->Status);
+        resp.set_result_set_index(static_cast<i64>(GetProtoRequest()->result_set_index()));
+        if (ev->Get()->Issues) {
+            NYql::IssuesToMessage(ev->Get()->Issues, resp.mutable_issues());
+        }
+        if (ev->Get()->ResultSet) {
+            resp.mutable_result_set()->Swap(&(*ev->Get()->ResultSet));
+            if (ev->Get()->HasMoreResults) {
+                resp.set_next_fetch_token(ToString(RowsOffset + resp.result_set().rows_size()));
             }
         }
         Reply(resp.status(), std::move(resp));

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

@@ -142,9 +142,6 @@ struct TEvKqp {
 
     using TEvAbortExecution = NYql::NDq::TEvDq::TEvAbortExecution;
 
-    struct TEvFetchScriptResultsResponse : public TEventPB<TEvFetchScriptResultsResponse, NKikimrKqp::TEvFetchScriptResultsResponse, TKqpEvents::EvFetchScriptResultsResponse> {
-    };
-
     struct TEvCancelScriptExecutionRequest : public TEventPB<TEvCancelScriptExecutionRequest, NKikimrKqp::TEvCancelScriptExecutionRequest, TKqpEvents::EvCancelScriptExecutionRequest> {
     };
 

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

@@ -224,15 +224,18 @@ struct TEvSaveScriptResultFinished : public NActors::TEventLocal<TEvSaveScriptRe
     NYql::TIssues Issues;
 };
 
-struct TEvFetchScriptResultsQueryResponse : public NActors::TEventLocal<TEvFetchScriptResultsQueryResponse, TKqpScriptExecutionEvents::EvFetchScriptResultsQueryResponse> {
-    TEvFetchScriptResultsQueryResponse(bool truncated, NKikimrKqp::TEvFetchScriptResultsResponse&& results)
-        : Truncated(truncated)
-        , Results(std::move(results))
-    {
-    }
+struct TEvFetchScriptResultsResponse : public NActors::TEventLocal<TEvFetchScriptResultsResponse, TKqpScriptExecutionEvents::EvFetchScriptResultsResponse> {
+    TEvFetchScriptResultsResponse(Ydb::StatusIds::StatusCode status, std::optional<Ydb::ResultSet>&& resultSet, bool hasMoreResults, NYql::TIssues issues)
+        : Status(status)
+        , ResultSet(std::move(resultSet))
+        , HasMoreResults(hasMoreResults)
+        , Issues(std::move(issues))
+    {}
 
-    bool Truncated;
-    NKikimrKqp::TEvFetchScriptResultsResponse Results;
+    Ydb::StatusIds::StatusCode Status;
+    std::optional<Ydb::ResultSet> ResultSet;
+    bool HasMoreResults;
+    NYql::TIssues Issues;
 };
 
 struct TEvSaveScriptExternalEffectRequest : public NActors::TEventLocal<TEvSaveScriptExternalEffectRequest, TKqpScriptExecutionEvents::EvSaveScriptExternalEffectRequest> {

ydb/core/kqp/common/simple/kqp_event_ids.h

@@ -148,7 +148,7 @@ struct TKqpScriptExecutionEvents {
         EvSaveScriptResultFinished,
         EvCheckAliveRequest,
         EvCheckAliveResponse,
-        EvFetchScriptResultsQueryResponse,
+        EvFetchScriptResultsResponse,
         EvSaveScriptExternalEffectRequest,
         EvSaveScriptExternalEffectResponse,
         EvScriptFinalizeRequest,