Skip to content

Commit 148e7bc

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

File tree

10 files changed

+225
-86
lines changed

10 files changed

+225
-86
lines changed

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

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ using TEvExecuteQueryRequest = TGrpcRequestNoOperationCall<Ydb::Query::ExecuteQu
2424
struct TProducerState {
2525
TMaybe<ui64> LastSeqNo;
2626
ui64 AckedFreeSpaceBytes = 0;
27+
TActorId ActorId;
2728
};
2829

2930
class TRpcFlowControlState {
@@ -335,8 +336,9 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
335336
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
336337
resp->Record.SetSeqNo(*producer.LastSeqNo);
337338
resp->Record.SetFreeSpace(freeSpaceBytes);
339+
resp->Record.SetQueryResultIndex(producerId);
338340

339-
ctx.Send(producerId, resp.Release());
341+
ctx.Send(producer.ActorId, resp.Release());
340342

341343
producer.AckedFreeSpaceBytes = freeSpaceBytes;
342344
}
@@ -347,6 +349,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
347349
void Handle(NKqp::TEvKqpExecuter::TEvStreamData::TPtr& ev, const TActorContext& ctx) {
348350
Ydb::Query::ExecuteQueryResponsePart response;
349351
response.set_status(Ydb::StatusIds::SUCCESS);
352+
ui64 queryResultIndex = ev->Get()->Record.GetQueryResultIndex();
350353
response.set_result_set_index(ev->Get()->Record.GetQueryResultIndex());
351354
response.mutable_result_set()->Swap(ev->Get()->Record.MutableResultSet());
352355

@@ -358,7 +361,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
358361

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

361-
auto& producer = StreamProducers_[ev->Sender];
364+
auto& producer = StreamProducers_[queryResultIndex];
365+
producer.ActorId = ev->Sender;
362366
producer.LastSeqNo = ev->Get()->Record.GetSeqNo();
363367
producer.AckedFreeSpaceBytes = freeSpaceBytes;
364368

@@ -371,6 +375,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
371375
auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
372376
resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
373377
resp->Record.SetFreeSpace(freeSpaceBytes);
378+
resp->Record.SetQueryResultIndex(queryResultIndex);
374379

375380
ctx.Send(ev->Sender, resp.Release());
376381
}
@@ -381,14 +386,27 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
381386
const auto& issueMessage = record.GetResponse().GetQueryIssues();
382387

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

388401
auto& kqpResponse = record.GetResponse();
389402

390403
Ydb::Query::ExecuteQueryResponsePart response;
391404

405+
for(int i = 0; i < kqpResponse.GetYdbResults().size(); i++) {
406+
response.set_result_set_index(i);
407+
response.mutable_result_set()->Swap(record.MutableResponse()->MutableYdbResults(i));
408+
}
409+
392410
AuditContextAppend(Request_.get(), *Request_->GetProtoRequest(), response);
393411

394412
if (kqpResponse.HasTxMeta()) {
@@ -493,7 +511,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
493511
std::shared_ptr<TEvExecuteQueryRequest> Request_;
494512

495513
TRpcFlowControlState FlowControl_;
496-
TMap<TActorId, TProducerState> StreamProducers_;
514+
TMap<ui64, TProducerState> StreamProducers_;
497515
};
498516

499517
} // namespace

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

Lines changed: 5 additions & 43 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()) {
@@ -1770,10 +1736,6 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
17701736
THashMap<ui64, TVector<ui64>> remoteComputeTasks; // shardId -> [task]
17711737
TVector<ui64> computeTasks;
17721738

1773-
if (StreamResult) {
1774-
InitializeChannelProxies();
1775-
}
1776-
17771739
for (auto& task : TasksGraph.GetTasks()) {
17781740
auto& stageInfo = TasksGraph.GetStageInfo(task.StageId);
17791741
if (task.Meta.ShardId && (task.Meta.Reads || task.Meta.Writes)) {

‎ydb/core/kqp/executer_actor/kqp_executer_impl.h‎

Lines changed: 132 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
#include <ydb/core/kqp/common/kqp_ru_calc.h>
1313
#include <ydb/core/kqp/common/kqp_lwtrace_probes.h>
14+
#include <ydb/core/kqp/runtime/kqp_transport.h>
1415

1516
#include <ydb/core/actorlib_impl/long_timer.h>
1617
#include <ydb/core/base/appdata.h>
@@ -34,6 +35,7 @@
3435

3536
#include <ydb/library/yql/dq/actors/compute/dq_compute_actor.h>
3637
#include <ydb/library/yql/dq/runtime/dq_transport.h>
38+
#include <ydb/library/yql/dq/common/dq_serialized_batch.h>
3739
#include <ydb/library/yql/providers/common/http_gateway/yql_http_gateway.h>
3840
#include <ydb/library/yql/providers/common/structured_token/yql_token_builder.h>
3941
#include <ydb/library/yql/public/issue/yql_issue.h>
@@ -45,6 +47,7 @@
4547
#include <ydb/library/actors/core/hfunc.h>
4648
#include <ydb/library/actors/core/log.h>
4749

50+
4851
#include <util/generic/size_literals.h>
4952

5053

@@ -119,7 +122,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
119122
const NKikimrConfig::TTableServiceConfig::EChannelTransportVersion chanTransportVersion,
120123
const NKikimrConfig::TTableServiceConfig::TAggregationConfig& aggregation,
121124
TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
122-
ui64 spanVerbosity = 0, TString spanName = "KqpExecuterBase")
125+
ui64 spanVerbosity = 0, TString spanName = "KqpExecuterBase", bool streamResult = false)
123126
: Request(std::move(request))
124127
, Database(database)
125128
, UserToken(userToken)
@@ -130,6 +133,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
130133
, MaximalSecretsSnapshotWaitTime(maximalSecretsSnapshotWaitTime)
131134
, AggregationSettings(aggregation)
132135
, HasOlapTable(false)
136+
, StreamResult(streamResult)
133137
{
134138
TasksGraph.GetMeta().Snapshot = IKqpGateway::TKqpSnapshot(Request.Snapshot.Step, Request.Snapshot.TxId);
135139
TasksGraph.GetMeta().Arena = MakeIntrusive<NActors::TProtoArenaHolder>();
@@ -234,6 +238,130 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
234238
return true;
235239
}
236240

241+
struct TEvComputeChannelDataOOB {
242+
NYql::NDqProto::TEvComputeChannelData Proto;
243+
TRope Payload;
244+
245+
size_t Size() const {
246+
return Proto.GetChannelData().GetData().GetRaw().size() + Payload.size();
247+
}
248+
249+
ui32 RowCount() const {
250+
return Proto.GetChannelData().GetData().GetRows();
251+
}
252+
};
253+
254+
void HandleChannelData(NYql::NDq::TEvDqCompute::TEvChannelData::TPtr& ev) {
255+
auto& record = ev->Get()->Record;
256+
auto& channelData = record.GetChannelData();
257+
auto& channel = TasksGraph.GetChannel(channelData.GetChannelId());
258+
auto& task = TasksGraph.GetTask(channel.SrcTask);
259+
const TActorId channelComputeActorId = ev->Sender;
260+
261+
auto& txResult = ResponseEv->TxResults[channel.DstInputIndex];
262+
ui32 queryResultIndex = txResult.QueryResultIndex.Defined() ? *txResult.QueryResultIndex : 0;
263+
auto [it, _] = ResultChannelToComputeActor.emplace(queryResultIndex, std::make_pair(ev->Sender, channel.Id));
264+
YQL_ENSURE(it->second.first == channelComputeActorId);
265+
266+
if (StreamResult && txResult.IsStream && txResult.QueryResultIndex.Defined()) {
267+
268+
TEvComputeChannelDataOOB computeData;
269+
computeData.Proto = std::move(ev->Get()->Record);
270+
if (computeData.Proto.GetChannelData().GetData().HasPayloadId()) {
271+
computeData.Payload = ev->Get()->GetPayload(computeData.Proto.GetChannelData().GetData().GetPayloadId());
272+
}
273+
274+
const bool trailingResults = (
275+
computeData.Proto.GetChannelData().GetFinished() &&
276+
(ResponseEv->TxResults.size() == 1) &&
277+
Request.IsTrailingResultsAllowed());
278+
279+
TVector<NYql::NDq::TDqSerializedBatch> batches(1);
280+
auto& batch = batches.front();
281+
282+
batch.Proto = std::move(*computeData.Proto.MutableChannelData()->MutableData());
283+
batch.Payload = std::move(computeData.Payload);
284+
285+
TKqpProtoBuilder protoBuilder{*AppData()->FunctionRegistry};
286+
auto resultSet = protoBuilder.BuildYdbResultSet(std::move(batches), txResult.MkqlItemType, txResult.ColumnOrder);
287+
288+
if (!trailingResults) {
289+
auto streamEv = MakeHolder<TEvKqpExecuter::TEvStreamData>();
290+
streamEv->Record.SetSeqNo(computeData.Proto.GetSeqNo());
291+
streamEv->Record.SetQueryResultIndex(*txResult.QueryResultIndex);
292+
streamEv->Record.MutableResultSet()->Swap(&resultSet);
293+
294+
LOG_D("Send TEvStreamData to " << Target << ", seqNo: " << streamEv->Record.GetSeqNo()
295+
<< ", nRows: " << streamEv->Record.GetResultSet().rows().size());
296+
297+
this->Send(Target, streamEv.Release());
298+
299+
} else {
300+
auto ackEv = MakeHolder<NYql::NDq::TEvDqCompute::TEvChannelDataAck>();
301+
ackEv->Record.SetSeqNo(computeData.Proto.GetSeqNo());
302+
ackEv->Record.SetChannelId(channel.Id);
303+
ackEv->Record.SetFreeSpace(50_MB);
304+
this->Send(channelComputeActorId, ackEv.Release(), /* TODO: undelivery */ 0, /* cookie */ channel.Id);
305+
txResult.TrailingResult.Swap(&resultSet);
306+
txResult.HasTrailingResult = true;
307+
LOG_D("staging TEvStreamData to " << Target << ", seqNo: " << computeData.Proto.GetSeqNo()
308+
<< ", nRows: " << txResult.TrailingResult.rows().size());
309+
}
310+
311+
return;
312+
}
313+
314+
NYql::NDq::TDqSerializedBatch batch;
315+
batch.Proto = std::move(*record.MutableChannelData()->MutableData());
316+
if (batch.Proto.HasPayloadId()) {
317+
batch.Payload = ev->Get()->GetPayload(batch.Proto.GetPayloadId());
318+
}
319+
320+
YQL_ENSURE(channel.DstTask == 0);
321+
322+
if (Stats) {
323+
Stats->ResultBytes += batch.Size();
324+
Stats->ResultRows += batch.RowCount();
325+
}
326+
327+
LOG_T("Got result, channelId: " << channel.Id << ", shardId: " << task.Meta.ShardId
328+
<< ", inputIndex: " << channel.DstInputIndex << ", from: " << ev->Sender
329+
<< ", finished: " << channelData.GetFinished());
330+
331+
ResponseEv->TakeResult(channel.DstInputIndex, std::move(batch));
332+
LOG_T("Send ack to channelId: " << channel.Id << ", seqNo: " << record.GetSeqNo() << ", to: " << ev->Sender);
333+
334+
auto ackEv = MakeHolder<NYql::NDq::TEvDqCompute::TEvChannelDataAck>();
335+
ackEv->Record.SetSeqNo(record.GetSeqNo());
336+
ackEv->Record.SetChannelId(channel.Id);
337+
ackEv->Record.SetFreeSpace(50_MB);
338+
this->Send(channelComputeActorId, ackEv.Release(), /* TODO: undelivery */ 0, /* cookie */ channel.Id);
339+
}
340+
341+
void HandleStreamAck(TEvKqpExecuter::TEvStreamDataAck::TPtr& ev) {
342+
ui64 queryResultIndex = ev->Get()->Record.GetQueryResultIndex();
343+
auto it = ResultChannelToComputeActor.find(queryResultIndex);
344+
YQL_ENSURE(it != ResultChannelToComputeActor.end());
345+
const auto [channelComputeActorId, channelId] = it->second;
346+
347+
ui64 seqNo = ev->Get()->Record.GetSeqNo();
348+
i64 freeSpace = ev->Get()->Record.GetFreeSpace();
349+
350+
LOG_ERROR_S(*NActors::TlsActivationContext, NKikimrServices::KQP_EXECUTER, "TxId: " << TxId
351+
<< ", send ack to channelId: " << channelId
352+
<< ", seqNo: " << seqNo
353+
<< ", enough: " << ev->Get()->Record.GetEnough()
354+
<< ", freeSpace: " << freeSpace
355+
<< ", to: " << channelComputeActorId);
356+
357+
auto ackEv = MakeHolder<NYql::NDq::TEvDqCompute::TEvChannelDataAck>();
358+
ackEv->Record.SetSeqNo(seqNo);
359+
ackEv->Record.SetChannelId(channelId);
360+
ackEv->Record.SetFreeSpace(freeSpace);
361+
ackEv->Record.SetFinish(ev->Get()->Record.GetEnough());
362+
this->Send(channelComputeActorId, ackEv.Release(), /* TODO: undelivery */ 0, /* cookie */ channelId);
363+
}
364+
237365
void HandleComputeStats(NYql::NDq::TEvDqCompute::TEvState::TPtr& ev) {
238366
TActorId computeActor = ev->Sender;
239367
auto& state = ev->Get()->Record;
@@ -1594,16 +1722,6 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
15941722
return true;
15951723
}
15961724

1597-
void InitializeChannelProxies() {
1598-
for(const auto& channel: TasksGraph.GetChannels()) {
1599-
if (channel.DstTask) {
1600-
continue;
1601-
}
1602-
1603-
CreateChannelProxy(channel);
1604-
}
1605-
}
1606-
16071725
const IKqpGateway::TKqpSnapshot& GetSnapshot() const {
16081726
return TasksGraph.GetMeta().Snapshot;
16091727
}
@@ -1753,8 +1871,11 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
17531871
const NKikimrConfig::TTableServiceConfig::TAggregationConfig AggregationSettings;
17541872
TVector<NKikimrKqp::TKqpNodeResources> ResourcesSnapshot;
17551873
bool HasOlapTable = false;
1874+
bool StreamResult = false;
17561875
bool HasDatashardSourceScan = false;
17571876
bool UnknownAffectedShardCount = false;
1877+
1878+
THashMap<ui64, std::pair<TActorId, ui64>> ResultChannelToComputeActor;
17581879
THashMap<NYql::NDq::TStageId, THashMap<ui64, TShardInfo>> SourceScanStageIdToParititions;
17591880

17601881
private:

0 commit comments

Comments
 (0)