Skip to content

Commit df7bfaf

Browse files
shnikdshnikd
committed
Initial commit
1 parent b4f0145 commit df7bfaf

File tree

3 files changed

+147
-47
lines changed

3 files changed

+147
-47
lines changed

‎ydb/core/kqp/compile_service/kqp_compile_service.cpp

+68-47
Original file line numberDiff line numberDiff line change
@@ -566,24 +566,26 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
566566
if (request.Uid) {
567567
Counters->ReportCompileRequestGet(dbCounters);
568568

569-
auto compileResult = QueryCache.FindByUid(*request.Uid, request.KeepInCache);
570-
if (compileResult) {
571-
Y_ENSURE(compileResult->Query);
572-
if (compileResult->Query->UserSid == userSid) {
573-
Counters->ReportQueryCacheHit(dbCounters, true);
574-
575-
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache by uid"
576-
<< ", sender: " << ev->Sender
577-
<< ", queryUid: " << *request.Uid);
578-
579-
ReplyFromCache(ev->Sender, compileResult, ctx, ev->Cookie, std::move(ev->Get()->Orbit), std::move(compileServiceSpan));
580-
return;
581-
} else {
582-
LOG_NOTICE_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Non-matching user sid for query"
583-
<< ", sender: " << ev->Sender
584-
<< ", queryUid: " << *request.Uid
585-
<< ", expected sid: " << compileResult->Query->UserSid
586-
<< ", actual sid: " << userSid);
569+
if (!request.TempTablesState || request.TempTablesState->TempTables.empty()) {
570+
auto compileResult = QueryCache.FindByUid(*request.Uid, request.KeepInCache);
571+
if (compileResult) {
572+
Y_ENSURE(compileResult->Query);
573+
if (compileResult->Query->UserSid == userSid) {
574+
Counters->ReportQueryCacheHit(dbCounters, true);
575+
576+
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache by uid"
577+
<< ", sender: " << ev->Sender
578+
<< ", queryUid: " << *request.Uid);
579+
580+
ReplyFromCache(ev->Sender, compileResult, ctx, ev->Cookie, std::move(ev->Get()->Orbit), std::move(compileServiceSpan));
581+
return;
582+
} else {
583+
LOG_NOTICE_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Non-matching user sid for query"
584+
<< ", sender: " << ev->Sender
585+
<< ", queryUid: " << *request.Uid
586+
<< ", expected sid: " << compileResult->Query->UserSid
587+
<< ", actual sid: " << userSid);
588+
}
587589
}
588590
}
589591

@@ -609,16 +611,19 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
609611
Y_ENSURE(query.UserSid == userSid);
610612
}
611613

612-
auto compileResult = QueryCache.FindByQuery(query, request.KeepInCache);
613-
if (compileResult) {
614-
Counters->ReportQueryCacheHit(dbCounters, true);
615614

616-
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache from query text"
617-
<< ", sender: " << ev->Sender
618-
<< ", queryUid: " << compileResult->Uid);
615+
if (!request.TempTablesState || request.TempTablesState->TempTables.empty()) {
616+
auto compileResult = QueryCache.FindByQuery(query, request.KeepInCache);
617+
if (compileResult) {
618+
Counters->ReportQueryCacheHit(dbCounters, true);
619619

620-
ReplyFromCache(ev->Sender, compileResult, ctx, ev->Cookie, std::move(ev->Get()->Orbit), std::move(compileServiceSpan));
621-
return;
620+
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache from query text"
621+
<< ", sender: " << ev->Sender
622+
<< ", queryUid: " << compileResult->Uid);
623+
624+
ReplyFromCache(ev->Sender, compileResult, ctx, ev->Cookie, std::move(ev->Get()->Orbit), std::move(compileServiceSpan));
625+
return;
626+
}
622627
}
623628

624629
CollectDiagnostics = request.CollectDiagnostics;
@@ -672,7 +677,11 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
672677
auto dbCounters = request.DbCounters;
673678
Counters->ReportRecompileRequestGet(dbCounters);
674679

675-
auto compileResult = QueryCache.FindByUid(request.Uid, false);
680+
TKqpCompileResult::TConstPtr compileResult = nullptr;
681+
if (!request.TempTablesState || request.TempTablesState->TempTables.empty()) {
682+
compileResult = QueryCache.FindByUid(request.Uid, false);
683+
}
684+
676685
if (compileResult || request.Query) {
677686
Counters->ReportCompileRequestCompile(dbCounters);
678687

@@ -736,18 +745,26 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
736745

737746
bool keepInCache = compileRequest.KeepInCache && compileResult->AllowCache;
738747

748+
bool hasTempTables = compileRequest.TempTablesState
749+
&& (!compileRequest.TempTablesState->TempTables.empty());
750+
if (compileResult->PreparedQuery) {
751+
hasTempTables = compileResult->PreparedQuery->HasTempTables(compileRequest.TempTablesState);
752+
}
753+
739754
try {
740755
if (compileResult->Status == Ydb::StatusIds::SUCCESS) {
741-
if (QueryCache.FindByUid(compileResult->Uid, false)) {
742-
QueryCache.Replace(compileResult);
743-
} else if (keepInCache) {
744-
if (QueryCache.Insert(compileResult, TableServiceConfig.GetEnableAstCache())) {
745-
Counters->CompileQueryCacheEvicted->Inc();
746-
}
747-
if (compileResult->Query && compileResult->Query->Settings.IsPrepareQuery) {
748-
if (InsertPreparingQuery(compileResult, compileRequest.KeepInCache)) {
756+
if (!hasTempTables) {
757+
if (QueryCache.FindByUid(compileResult->Uid, false)) {
758+
QueryCache.Replace(compileResult);
759+
} else if (keepInCache) {
760+
if (QueryCache.Insert(compileResult, TableServiceConfig.GetEnableAstCache())) {
749761
Counters->CompileQueryCacheEvicted->Inc();
750-
};
762+
}
763+
if (compileResult->Query && compileResult->Query->Settings.IsPrepareQuery) {
764+
if (InsertPreparingQuery(compileResult, compileRequest.KeepInCache)) {
765+
Counters->CompileQueryCacheEvicted->Inc();
766+
};
767+
}
751768
}
752769
}
753770

@@ -762,8 +779,10 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
762779
request.Cookie, std::move(request.Orbit), std::move(request.CompileServiceSpan), (CollectDiagnostics ? ev->Get()->ReplayMessageUserView : std::nullopt));
763780
}
764781
} else {
765-
if (QueryCache.FindByUid(compileResult->Uid, false)) {
766-
QueryCache.EraseByUid(compileResult->Uid);
782+
if (!hasTempTables) {
783+
if (QueryCache.FindByUid(compileResult->Uid, false)) {
784+
QueryCache.EraseByUid(compileResult->Uid);
785+
}
767786
}
768787
}
769788

@@ -819,18 +838,20 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
819838
auto& query = ev->Get()->Query;
820839
auto compileRequest = RequestsQueue.FinishActiveRequest(query);
821840
if (parseResult && parseResult->Ast->IsOk()) {
822-
auto compileResult = QueryCache.FindByAst(query, *parseResult->Ast, compileRequest.KeepInCache);
823-
if (compileResult) {
824-
Counters->ReportQueryCacheHit(compileRequest.DbCounters, true);
841+
if (!compileRequest.TempTablesState || compileRequest.TempTablesState->TempTables.empty()) {
842+
auto compileResult = QueryCache.FindByAst(query, *parseResult->Ast, compileRequest.KeepInCache);
843+
if (compileResult) {
844+
Counters->ReportQueryCacheHit(compileRequest.DbCounters, true);
825845

826-
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache from ast"
827-
<< ", sender: " << compileRequest.Sender
828-
<< ", queryUid: " << compileResult->Uid);
846+
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Served query from cache from ast"
847+
<< ", sender: " << compileRequest.Sender
848+
<< ", queryUid: " << compileResult->Uid);
829849

830-
compileResult->Ast->PgAutoParamValues = std::move(parseResult->Ast->PgAutoParamValues);
850+
compileResult->Ast->PgAutoParamValues = std::move(parseResult->Ast->PgAutoParamValues);
831851

832-
ReplyFromCache(compileRequest.Sender, compileResult, ctx, compileRequest.Cookie, std::move(compileRequest.Orbit), std::move(compileRequest.CompileServiceSpan));
833-
return;
852+
ReplyFromCache(compileRequest.Sender, compileResult, ctx, compileRequest.Cookie, std::move(compileRequest.Orbit), std::move(compileRequest.CompileServiceSpan));
853+
return;
854+
}
834855
}
835856
}
836857
Counters->ReportQueryCacheHit(compileRequest.DbCounters, false);

‎ydb/core/kqp/query_data/kqp_prepared_query.cpp

+73
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "kqp_prepared_query.h"
22

3+
#include <ydb/core/base/path.h>
34
#include <ydb/core/kqp/common/kqp_resolve.h>
45
#include <ydb/library/mkql_proto/mkql_proto.h>
56
#include <ydb/core/kqp/provider/yql_kikimr_settings.h>
@@ -106,6 +107,49 @@ bool TKqpPhyTxHolder::IsLiteralTx() const {
106107
return LiteralTx;
107108
}
108109

110+
std::optional<std::pair<std::pair<TString, TString>, bool>>
111+
TKqpPhyTxHolder::GetSchemeOpTempTablePath() const {
112+
if (GetType() != NKqpProto::TKqpPhyTx::TYPE_SCHEME) {
113+
return std::nullopt;
114+
}
115+
auto& schemeOperation = GetSchemeOperation();
116+
switch (schemeOperation.GetOperationCase()) {
117+
case NKqpProto::TKqpSchemeOperation::kCreateTable: {
118+
const auto& modifyScheme = schemeOperation.GetCreateTable();
119+
const NKikimrSchemeOp::TTableDescription* tableDesc = nullptr;
120+
switch (modifyScheme.GetOperationType()) {
121+
case NKikimrSchemeOp::ESchemeOpCreateTable: {
122+
tableDesc = &modifyScheme.GetCreateTable();
123+
break;
124+
}
125+
case NKikimrSchemeOp::ESchemeOpCreateIndexedTable: {
126+
tableDesc = &modifyScheme.GetCreateIndexedTable().GetTableDescription();
127+
break;
128+
}
129+
default:
130+
return std::nullopt;
131+
}
132+
if (tableDesc->HasTemporary()) {
133+
if (tableDesc->GetTemporary()) {
134+
return {{{modifyScheme.GetWorkingDir(), tableDesc->GetName()},
135+
true}};
136+
}
137+
}
138+
break;
139+
}
140+
case NKqpProto::TKqpSchemeOperation::kDropTable: {
141+
auto modifyScheme = schemeOperation.GetDropTable();
142+
auto* dropTable = modifyScheme.MutableDrop();
143+
144+
return {{{modifyScheme.GetWorkingDir(), dropTable->GetName()},
145+
false}};
146+
}
147+
default:
148+
return std::nullopt;
149+
}
150+
return std::nullopt;
151+
}
152+
109153
const NKikimr::NKqp::TStagePredictor& TKqpPhyTxHolder::GetCalculationPredictor(const size_t stageIdx) const {
110154
YQL_ENSURE(stageIdx < Predictors.size(), "incorrect stage idx for predictor");
111155
return Predictors[stageIdx];
@@ -226,6 +270,35 @@ void TPreparedQueryHolder::FillTables(const google::protobuf::RepeatedPtrField<
226270
}
227271
}
228272

273+
bool TPreparedQueryHolder::HasTempTables(TKqpTempTablesState::TConstPtr tempTablesState) const {
274+
auto tempTables = THashSet<TString>();
275+
for (const auto& [path, info] : tempTablesState->TempTables) {
276+
tempTables.insert(path.second + *tempTablesState->SessionId);
277+
}
278+
for (const auto& table: QueryTables) {
279+
if (tempTables.contains(table)) {
280+
return true;
281+
}
282+
}
283+
284+
for (const auto& tx: Transactions) {
285+
auto optPath = tx->GetSchemeOpTempTablePath();
286+
if (!optPath) {
287+
continue;
288+
} else {
289+
const auto& [path, isCreate] = *optPath;
290+
if (isCreate) {
291+
return true;
292+
} else {
293+
if (tempTables.contains(JoinPath({path.first, path.second}))) {
294+
return true;
295+
}
296+
}
297+
}
298+
}
299+
return false;
300+
}
301+
229302
const TKqpPhyTxHolder::TConstPtr& TPreparedQueryHolder::GetPhyTx(ui32 txId) const {
230303
YQL_ENSURE(txId < Transactions.size());
231304
return Transactions[txId];

‎ydb/core/kqp/query_data/kqp_prepared_query.h

+6
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#pragma once
22

3+
#include <ydb/core/kqp/common/simple/temp_tables.h>
34
#include <ydb/core/kqp/query_data/kqp_predictor.h>
45
#include <ydb/core/scheme/scheme_tabledefs.h>
56
#include <ydb/core/protos/kqp.pb.h>
@@ -115,6 +116,9 @@ class TKqpPhyTxHolder {
115116
const std::shared_ptr<TPreparedQueryAllocHolder>& alloc, TIntrusivePtr<TTableConstInfoMap> tableConstInfoById);
116117

117118
bool IsLiteralTx() const;
119+
120+
std::optional<std::pair<std::pair<TString, TString>, bool>>
121+
GetSchemeOpTempTablePath() const;
118122
};
119123

120124
class TLlvmSettings {
@@ -187,6 +191,8 @@ class TPreparedQueryHolder {
187191
void FillTable(const NKqpProto::TKqpPhyTable& phyTable);
188192

189193
void FillTables(const google::protobuf::RepeatedPtrField< ::NKqpProto::TKqpPhyStage>& stages);
194+
195+
bool HasTempTables(TKqpTempTablesState::TConstPtr tempTablesState) const;
190196
};
191197

192198

0 commit comments

Comments
 (0)