Skip to content

Commit 0e90a02

Browse files
gridnevvvitgridnevvvit
authored
Merge aeab9d4 into cb5a70d 
2 parents cb5a70d + aeab9d4 commit 0e90a02

13 files changed

+261
-76
lines changed

‎ydb/core/grpc_services/query/rpc_execute_query.cpp

+26-4
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ using TEvExecuteQueryRequest = TGrpcRequestNoOperationCall<Ydb::Query::ExecuteQu
2424
struct TProducerState {
2525
TMaybe<ui64> LastSeqNo;
2626
ui64 AckedFreeSpaceBytes = 0;
27+
TActorId ActorId;
28+
ui64 QueryResultIndex = 0;
2729
};
2830

2931
class TRpcFlowControlState {
@@ -335,8 +337,10 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
335337
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
336338
resp->Record.SetSeqNo(*producer.LastSeqNo);
337339
resp->Record.SetFreeSpace(freeSpaceBytes);
340+
resp->Record.SetQueryResultIndex(producer.QueryResultIndex);
341+
resp->Record.SetChannelId(producerId);
338342

339-
ctx.Send(producerId, resp.Release());
343+
ctx.Send(producer.ActorId, resp.Release());
340344

341345
producer.AckedFreeSpaceBytes = freeSpaceBytes;
342346
}
@@ -358,9 +362,11 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
358362

359363
Request_->SendSerializedResult(std::move(out), Ydb::StatusIds::SUCCESS);
360364

361-
auto& producer = StreamProducers_[ev->Sender];
365+
auto& producer = StreamProducers_[ev->Get()->Record.GetChannelId()];
366+
producer.ActorId = ev->Sender;
362367
producer.LastSeqNo = ev->Get()->Record.GetSeqNo();
363368
producer.AckedFreeSpaceBytes = freeSpaceBytes;
369+
producer.QueryResultIndex = ev->Get()->Record.GetQueryResultIndex();
364370

365371
LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Send stream data ack"
366372
<< ", seqNo: " << ev->Get()->Record.GetSeqNo()
@@ -371,6 +377,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
371377
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
372378
resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
373379
resp->Record.SetFreeSpace(freeSpaceBytes);
380+
resp->Record.SetQueryResultIndex(ev->Get()->Record.GetQueryResultIndex());
381+
resp->Record.SetChannelId(ev->Get()->Record.GetChannelId());
374382

375383
ctx.Send(ev->Sender, resp.Release());
376384
}
@@ -381,14 +389,28 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
381389
const auto& issueMessage = record.GetResponse().GetQueryIssues();
382390

383391
bool hasTrailingMessage = false;
384-
392+
393+
auto& kqpResponse = record.GetResponse();
394+
if (kqpResponse.GetYdbResults().size() > 1) {
395+
auto issue = MakeIssue(NKikimrIssues::TIssuesIds::DEFAULT_ERROR,
396+
"Unexpected trailing message with multiple result sets.");
397+
ReplyFinishStream(Ydb::StatusIds::INTERNAL_ERROR, issue);
398+
return;
399+
}
400+
385401
if (record.GetYdbStatus() == Ydb::StatusIds::SUCCESS) {
386402
Request_->SetRuHeader(record.GetConsumedRu());
387403

388404
auto& kqpResponse = record.GetResponse();
389405

390406
Ydb::Query::ExecuteQueryResponsePart response;
391407

408+
for(int i = 0; i < kqpResponse.GetYdbResults().size(); i++) {
409+
hasTrailingMessage = true;
410+
response.set_result_set_index(i);
411+
response.mutable_result_set()->Swap(record.MutableResponse()->MutableYdbResults(i));
412+
}
413+
392414
AuditContextAppend(Request_.get(), *Request_->GetProtoRequest(), response);
393415

394416
if (kqpResponse.HasTxMeta()) {
@@ -497,7 +519,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
497519
std::shared_ptr<TEvExecuteQueryRequest> Request_;
498520

499521
TRpcFlowControlState FlowControl_;
500-
TMap<TActorId, TProducerState> StreamProducers_;
522+
TMap<ui64, TProducerState> StreamProducers_;
501523
};
502524

503525
} // namespace

‎ydb/core/grpc_services/rpc_stream_execute_scan_query.cpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,8 @@ class TStreamExecuteScanQueryRPC : public TActorBootstrapped<TStreamExecuteScanQ
295295
bool reportStats = NeedReportStats(*Request_->GetProtoRequest());
296296
bool reportPlan = reportStats && NeedReportPlan(*Request_->GetProtoRequest());
297297

298-
if (reportStats) {
298+
if (reportStats
299+
) {
299300
if (kqpResponse.HasQueryStats()) {
300301
for (const auto& execStats: ExecutionProfiles_) {
301302
record.MutableResponse()->MutableQueryStats()->AddExecutions()->Swap(execStats.get());

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

+7-1
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,8 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
3434
const ::Ydb::Operations::OperationParams* operationParams,
3535
bool keepSession = false,
3636
bool useCancelAfter = true,
37-
const ::Ydb::Query::Syntax syntax = Ydb::Query::Syntax::SYNTAX_UNSPECIFIED);
37+
const ::Ydb::Query::Syntax syntax = Ydb::Query::Syntax::SYNTAX_UNSPECIFIED,
38+
bool supportsStreamTrailingResult = false);
3839

3940
TEvQueryRequest() = default;
4041

@@ -285,6 +286,10 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
285286
ProgressStatsPeriod = progressStatsPeriod;
286287
}
287288

289+
bool GetSupportsStreamTrailingResult() const {
290+
return SupportsStreamTrailingResult;
291+
}
292+
288293
TDuration GetProgressStatsPeriod() const {
289294
return ProgressStatsPeriod;
290295
}
@@ -317,6 +322,7 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
317322
const ::Ydb::Query::Syntax Syntax = Ydb::Query::Syntax::SYNTAX_UNSPECIFIED;
318323
TIntrusivePtr<TUserRequestContext> UserRequestContext;
319324
TDuration ProgressStatsPeriod;
325+
bool SupportsStreamTrailingResult = false;
320326
};
321327

322328
struct TEvDataQueryStreamPart: public TEventPB<TEvDataQueryStreamPart,

‎ydb/core/kqp/common/kqp_event_impl.cpp

+3-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ TEvKqp::TEvQueryRequest::TEvQueryRequest(
2020
const ::Ydb::Operations::OperationParams* operationParams,
2121
bool keepSession,
2222
bool useCancelAfter,
23-
const ::Ydb::Query::Syntax syntax)
23+
const ::Ydb::Query::Syntax syntax,
24+
bool supportsStreamTrailingResult)
2425
: RequestCtx(ctx)
2526
, RequestActorId(requestActorId)
2627
, Database(CanonizePath(ctx->GetDatabaseName().GetOrElse("")))
@@ -36,6 +37,7 @@ TEvKqp::TEvQueryRequest::TEvQueryRequest(
3637
, HasOperationParams(operationParams)
3738
, KeepSession(keepSession)
3839
, Syntax(syntax)
40+
, SupportsStreamTrailingResult(supportsStreamTrailingResult)
3941
{
4042
if (HasOperationParams) {
4143
OperationTimeout = GetDuration(operationParams->operation_timeout());

‎ydb/core/kqp/executer_actor/kqp_data_executer.cpp

+5-40
Original file line numberDiff line numberDiff line change
@@ -130,10 +130,9 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
130130
const TActorId& creator, TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
131131
const bool enableOlapSink)
132132
: TBase(std::move(request), database, userToken, counters, executerRetriesConfig, chanTransportVersion, aggregation,
133-
maximalSecretsSnapshotWaitTime, userRequestContext, TWilsonKqp::DataExecuter, "DataExecuter"
133+
maximalSecretsSnapshotWaitTime, userRequestContext, TWilsonKqp::DataExecuter, "DataExecuter", streamResult
134134
)
135135
, AsyncIoFactory(std::move(asyncIoFactory))
136-
, StreamResult(streamResult)
137136
, EnableOlapSink(enableOlapSink)
138137
{
139138
Target = creator;
@@ -347,7 +346,8 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
347346
hFunc(TEvPersQueue::TEvProposeTransactionResult, HandlePrepare);
348347
hFunc(TEvPrivate::TEvReattachToShard, HandleExecute);
349348
hFunc(TEvDqCompute::TEvState, HandlePrepare); // from CA
350-
hFunc(TEvDqCompute::TEvChannelData, HandleExecute); // from CA
349+
hFunc(TEvDqCompute::TEvChannelData, HandleChannelData); // from CA
350+
hFunc(TEvKqpExecuter::TEvStreamDataAck, HandleStreamAck);
351351
hFunc(TEvPipeCache::TEvDeliveryProblem, HandlePrepare);
352352
hFunc(TEvKqp::TEvAbortExecution, HandlePrepare);
353353
hFunc(TEvents::TEvUndelivered, HandleUndelivered);
@@ -935,7 +935,8 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
935935
hFunc(TEvKqpNode::TEvStartKqpTasksResponse, HandleStartKqpTasksResponse);
936936
hFunc(TEvTxProxy::TEvProposeTransactionStatus, HandleExecute);
937937
hFunc(TEvDqCompute::TEvState, HandleComputeStats);
938-
hFunc(TEvDqCompute::TEvChannelData, HandleExecute);
938+
hFunc(NYql::NDq::TEvDqCompute::TEvChannelData, HandleChannelData);
939+
hFunc(TEvKqpExecuter::TEvStreamDataAck, HandleStreamAck);
939940
hFunc(TEvKqp::TEvAbortExecution, HandleExecute);
940941
IgnoreFunc(TEvInterconnect::TEvNodeConnected);
941942
default:
@@ -1286,41 +1287,6 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
12861287
}
12871288
}
12881289

1289-
void HandleExecute(TEvDqCompute::TEvChannelData::TPtr& ev) {
1290-
auto& record = ev->Get()->Record;
1291-
auto& channelData = record.GetChannelData();
1292-
1293-
TDqSerializedBatch batch;
1294-
batch.Proto = std::move(*record.MutableChannelData()->MutableData());
1295-
if (batch.Proto.HasPayloadId()) {
1296-
batch.Payload = ev->Get()->GetPayload(batch.Proto.GetPayloadId());
1297-
}
1298-
1299-
auto& channel = TasksGraph.GetChannel(channelData.GetChannelId());
1300-
YQL_ENSURE(channel.DstTask == 0);
1301-
auto shardId = TasksGraph.GetTask(channel.SrcTask).Meta.ShardId;
1302-
1303-
if (Stats) {
1304-
Stats->ResultBytes += batch.Size();
1305-
Stats->ResultRows += batch.RowCount();
1306-
}
1307-
1308-
LOG_T("Got result, channelId: " << channel.Id << ", shardId: " << shardId
1309-
<< ", inputIndex: " << channel.DstInputIndex << ", from: " << ev->Sender
1310-
<< ", finished: " << channelData.GetFinished());
1311-
1312-
ResponseEv->TakeResult(channel.DstInputIndex, std::move(batch));
1313-
{
1314-
LOG_T("Send ack to channelId: " << channel.Id << ", seqNo: " << record.GetSeqNo() << ", to: " << ev->Sender);
1315-
1316-
auto ackEv = MakeHolder<TEvDqCompute::TEvChannelDataAck>();
1317-
ackEv->Record.SetSeqNo(record.GetSeqNo());
1318-
ackEv->Record.SetChannelId(channel.Id);
1319-
ackEv->Record.SetFreeSpace(50_MB);
1320-
Send(ev->Sender, ackEv.Release(), /* TODO: undelivery */ 0, /* cookie */ channel.Id);
1321-
}
1322-
}
1323-
13241290
private:
13251291
bool IsReadOnlyTx() const {
13261292
if (Request.TopicOperations.HasOperations()) {
@@ -2417,7 +2383,6 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
24172383

24182384
private:
24192385
NYql::NDq::IDqAsyncIoFactory::TPtr AsyncIoFactory;
2420-
bool StreamResult = false;
24212386
bool EnableOlapSink = false;
24222387

24232388
bool HasExternalSources = false;

0 commit comments

Comments
 (0)