Skip to content

Commit e13b4f1

Browse files
gridnevvvitgridnevvvit
committed
support trailing generic query responses
1 parent 8ce7edb commit e13b4f1

File tree

12 files changed

+268
-91
lines changed

12 files changed

+268
-91
lines changed

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

Lines changed: 42 additions & 19 deletions
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 {
@@ -244,8 +246,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
244246
const auto traceId = Request_->GetTraceId();
245247

246248
NYql::TIssues issues;
247-
NKikimrKqp::EQueryAction queryAction;
248-
if (!ParseQueryAction(*req, queryAction, issues)) {
249+
if (!ParseQueryAction(*req, QueryAction, issues)) {
249250
return ReplyFinishStream(Ydb::StatusIds::BAD_REQUEST, std::move(issues));
250251
}
251252

@@ -274,7 +275,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
274275
cachePolicy->set_keep_in_cache(true);
275276

276277
auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(
277-
queryAction,
278+
QueryAction,
278279
queryType,
279280
SelfId(),
280281
Request_,
@@ -288,7 +289,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
288289
nullptr, // operationParams
289290
false, // keepSession
290291
false, // useCancelAfter
291-
syntax);
292+
syntax,
293+
true);
292294

293295
if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release())) {
294296
NYql::TIssues issues;
@@ -322,23 +324,24 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
322324

323325
ui64 freeSpaceBytes = FlowControl_.FreeSpaceBytes();
324326

325-
for (auto& pair : StreamProducers_) {
326-
const auto& producerId = pair.first;
327-
auto& producer = pair.second;
327+
for (auto& pair : StreamChannels_) {
328+
const auto& channelId = pair.first;
329+
auto& channel = pair.second;
328330

329-
if (freeSpaceBytes > 0 && producer.LastSeqNo && producer.AckedFreeSpaceBytes == 0) {
331+
if (freeSpaceBytes > 0 && channel.LastSeqNo && channel.AckedFreeSpaceBytes == 0) {
330332
LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Resume execution, "
331-
<< ", producer: " << producerId
332-
<< ", seqNo: " << producer.LastSeqNo
333+
<< ", channel: " << channelId
334+
<< ", seqNo: " << channel.LastSeqNo
333335
<< ", freeSpace: " << freeSpaceBytes);
334336

335337
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
336-
resp->Record.SetSeqNo(*producer.LastSeqNo);
338+
resp->Record.SetSeqNo(*channel.LastSeqNo);
337339
resp->Record.SetFreeSpace(freeSpaceBytes);
340+
resp->Record.SetChannelId(channelId);
338341

339-
ctx.Send(producerId, resp.Release());
342+
ctx.Send(channel.ActorId, resp.Release());
340343

341-
producer.AckedFreeSpaceBytes = freeSpaceBytes;
344+
channel.AckedFreeSpaceBytes = freeSpaceBytes;
342345
}
343346
}
344347

@@ -358,9 +361,11 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
358361

359362
Request_->SendSerializedResult(std::move(out), Ydb::StatusIds::SUCCESS);
360363

361-
auto& producer = StreamProducers_[ev->Sender];
362-
producer.LastSeqNo = ev->Get()->Record.GetSeqNo();
363-
producer.AckedFreeSpaceBytes = freeSpaceBytes;
364+
auto& channel = StreamChannels_[ev->Get()->Record.GetChannelId()];
365+
channel.ActorId = ev->Sender;
366+
channel.LastSeqNo = ev->Get()->Record.GetSeqNo();
367+
channel.AckedFreeSpaceBytes = freeSpaceBytes;
368+
channel.QueryResultIndex = ev->Get()->Record.GetQueryResultIndex();
364369

365370
LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Send stream data ack"
366371
<< ", seqNo: " << ev->Get()->Record.GetSeqNo()
@@ -371,8 +376,9 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
371376
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
372377
resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
373378
resp->Record.SetFreeSpace(freeSpaceBytes);
379+
resp->Record.SetChannelId(ev->Get()->Record.GetChannelId());
374380

375-
ctx.Send(ev->Sender, resp.Release());
381+
ctx.Send(channel.ActorId, resp.Release());
376382
}
377383

378384
void Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr& ev, const TActorContext&) {
@@ -381,14 +387,30 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
381387
const auto& issueMessage = record.GetResponse().GetQueryIssues();
382388

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

388402
auto& kqpResponse = record.GetResponse();
389403

390404
Ydb::Query::ExecuteQueryResponsePart response;
391405

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

394416
if (kqpResponse.HasTxMeta()) {
@@ -492,8 +514,9 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
492514
private:
493515
std::shared_ptr<TEvExecuteQueryRequest> Request_;
494516

517+
NKikimrKqp::EQueryAction QueryAction;
495518
TRpcFlowControlState FlowControl_;
496-
TMap<TActorId, TProducerState> StreamProducers_;
519+
TMap<ui64, TProducerState> StreamChannels_;
497520
};
498521

499522
} // namespace

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

Lines changed: 7 additions & 1 deletion
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‎

Lines changed: 3 additions & 1 deletion
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‎

Lines changed: 5 additions & 40 deletions
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)