Skip to content

Commit 7c678e3

Browse files
VPolkaVPolka
committed
recompilation with query id from compile result
1 parent ca21d47 commit 7c678e3

File tree

6 files changed

+38
-47
lines changed

6 files changed

+38
-47
lines changed

‎ydb/core/kqp/common/compilation/result.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,11 @@
22

33
namespace NKikimr::NKqp {
44

5+
std::shared_ptr<NYql::TAstParseResult> TKqpCompileResult::GetAst() const {
6+
if (QueryAst) {
7+
return QueryAst->Ast;
8+
}
9+
return nullptr;
10+
}
11+
512
} // namespace NKikimr::NKqp

‎ydb/core/kqp/common/compilation/result.h

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#pragma once
22
#include <memory>
3+
#include <ydb/core/kqp/common/simple/query_ast.h>
34
#include <ydb/core/kqp/common/simple/query_id.h>
45
#include <ydb/core/kqp/common/simple/helpers.h>
56
#include <ydb/library/yql/public/issue/yql_issue.h>
@@ -14,24 +15,26 @@ struct TKqpCompileResult {
1415
using TConstPtr = std::shared_ptr<const TKqpCompileResult>;
1516

1617
TKqpCompileResult(const TString& uid, const Ydb::StatusIds::StatusCode& status, const NYql::TIssues& issues,
17-
ETableReadType maxReadType, TMaybe<TKqpQueryId> query = {}, std::shared_ptr<NYql::TAstParseResult> ast = {},
18+
ETableReadType maxReadType, TMaybe<TKqpQueryId> query = {}, TMaybe<TQueryAst> queryAst = {},
1819
bool needToSplit = false, const TMaybe<TString>& commandTagName = {})
1920
: Status(status)
2021
, Issues(issues)
2122
, Query(std::move(query))
2223
, Uid(uid)
2324
, MaxReadType(maxReadType)
24-
, Ast(std::move(ast))
25+
, QueryAst(std::move(queryAst))
2526
, NeedToSplit(needToSplit)
2627
, CommandTagName(commandTagName) {}
2728

2829
static std::shared_ptr<TKqpCompileResult> Make(const TString& uid, const Ydb::StatusIds::StatusCode& status,
2930
const NYql::TIssues& issues, ETableReadType maxReadType, TMaybe<TKqpQueryId> query = {},
30-
std::shared_ptr<NYql::TAstParseResult> ast = {}, bool needToSplit = false, const TMaybe<TString>& commandTagName = {})
31+
TMaybe<TQueryAst> queryAst = {}, bool needToSplit = false, const TMaybe<TString>& commandTagName = {})
3132
{
32-
return std::make_shared<TKqpCompileResult>(uid, status, issues, maxReadType, std::move(query), std::move(ast), needToSplit, commandTagName);
33+
return std::make_shared<TKqpCompileResult>(uid, status, issues, maxReadType, std::move(query), std::move(queryAst), needToSplit, commandTagName);
3334
}
3435

36+
std::shared_ptr<NYql::TAstParseResult> GetAst() const;
37+
3538
Ydb::StatusIds::StatusCode Status;
3639
NYql::TIssues Issues;
3740

@@ -40,7 +43,7 @@ struct TKqpCompileResult {
4043

4144
ETableReadType MaxReadType;
4245
bool AllowCache = true;
43-
std::shared_ptr<NYql::TAstParseResult> Ast;
46+
TMaybe<TQueryAst> QueryAst;
4447
bool NeedToSplit = false;
4548
TMaybe<TString> CommandTagName = {};
4649

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

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -379,9 +379,9 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
379379

380380
void ReplyError(Ydb::StatusIds::StatusCode status, const TIssues& issues) {
381381
if (!KqpCompileResult) {
382-
KqpCompileResult = TKqpCompileResult::Make(Uid, status, issues, ETableReadType::Other, std::move(QueryId));
382+
KqpCompileResult = TKqpCompileResult::Make(Uid, status, issues, ETableReadType::Other, std::move(QueryId), std::move(QueryAst));
383383
} else {
384-
KqpCompileResult = TKqpCompileResult::Make(Uid, status, issues, ETableReadType::Other, std::move(KqpCompileResult->Query));
384+
KqpCompileResult = TKqpCompileResult::Make(Uid, status, issues, ETableReadType::Other, std::move(KqpCompileResult->Query), std::move(KqpCompileResult->QueryAst));
385385
}
386386

387387
Reply();
@@ -456,10 +456,6 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
456456
preparedQueryHolder->MutableLlvmSettings().Fill(Config, queryType);
457457
KqpCompileResult->PreparedQuery = preparedQueryHolder;
458458
KqpCompileResult->AllowCache = CanCacheQuery(KqpCompileResult->PreparedQuery->GetPhysicalQuery()) && allowCache;
459-
460-
if (QueryAst) {
461-
KqpCompileResult->Ast = QueryAst->Ast;
462-
}
463459
}
464460

465461
void Handle(TEvKqp::TEvContinueProcess::TPtr &ev, const TActorContext &ctx) {
@@ -478,7 +474,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
478474

479475
if (kqpResult.NeedToSplit) {
480476
KqpCompileResult = TKqpCompileResult::Make(
481-
Uid, status, kqpResult.Issues(), ETableReadType::Other, std::move(QueryId), {}, true);
477+
Uid, status, kqpResult.Issues(), ETableReadType::Other, std::move(QueryId), std::move(QueryAst), true);
482478
Reply();
483479
return;
484480
}
@@ -496,7 +492,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
496492

497493
auto queryType = QueryId.Settings.QueryType;
498494

499-
KqpCompileResult = TKqpCompileResult::Make(Uid, status, kqpResult.Issues(), maxReadType, std::move(QueryId));
495+
KqpCompileResult = TKqpCompileResult::Make(Uid, status, kqpResult.Issues(), maxReadType, std::move(QueryId), std::move(QueryAst));
500496
KqpCompileResult->CommandTagName = kqpResult.CommandTagName;
501497

502498
if (status == Ydb::StatusIds::SUCCESS) {

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

Lines changed: 15 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -50,16 +50,16 @@ class TKqpQueryCache {
5050

5151
void InsertAst(const TKqpCompileResult::TConstPtr& compileResult) {
5252
Y_ENSURE(compileResult->Query);
53-
Y_ENSURE(compileResult->Ast);
53+
Y_ENSURE(compileResult->GetAst());
5454

55-
AstIndex.emplace(GetQueryIdWithAst(*compileResult->Query, *compileResult->Ast), compileResult->Uid);
55+
AstIndex.emplace(GetQueryIdWithAst(*compileResult->Query, *compileResult->GetAst()), compileResult->Uid);
5656
}
5757

5858
bool Insert(const TKqpCompileResult::TConstPtr& compileResult, bool isEnableAstCache, bool isPerStatementExecution) {
5959
if (!isPerStatementExecution) {
6060
InsertQuery(compileResult);
6161
}
62-
if (isEnableAstCache && compileResult->Ast) {
62+
if (isEnableAstCache && compileResult->GetAst()) {
6363
InsertAst(compileResult);
6464
}
6565

@@ -76,8 +76,8 @@ class TKqpQueryCache {
7676

7777
auto queryId = *removedItem->Value.CompileResult->Query;
7878
QueryIndex.erase(queryId);
79-
if (removedItem->Value.CompileResult->Ast) {
80-
AstIndex.erase(GetQueryIdWithAst(queryId, *removedItem->Value.CompileResult->Ast));
79+
if (removedItem->Value.CompileResult->GetAst()) {
80+
AstIndex.erase(GetQueryIdWithAst(queryId, *removedItem->Value.CompileResult->GetAst()));
8181
}
8282
auto indexIt = Index.find(*removedItem);
8383
if (indexIt != Index.end()) {
@@ -190,8 +190,8 @@ class TKqpQueryCache {
190190
Y_ABORT_UNLESS(item->Value.CompileResult->Query);
191191
auto queryId = *item->Value.CompileResult->Query;
192192
QueryIndex.erase(queryId);
193-
if (item->Value.CompileResult->Ast) {
194-
AstIndex.erase(GetQueryIdWithAst(queryId, *item->Value.CompileResult->Ast));
193+
if (item->Value.CompileResult->GetAst()) {
194+
AstIndex.erase(GetQueryIdWithAst(queryId, *item->Value.CompileResult->GetAst()));
195195
}
196196

197197
Index.erase(it);
@@ -327,6 +327,8 @@ struct TKqpCompileRequest {
327327
NYql::TExprContext* SplitCtx;
328328
NYql::TExprNode::TPtr SplitExpr;
329329

330+
bool FindInCache = true;
331+
330332
bool IsIntrestedInResult() const {
331333
return IntrestedInResult->load();
332334
}
@@ -769,8 +771,6 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
769771
}
770772

771773
if (compileResult || request.Query) {
772-
QueryCache.EraseByUid(request.Uid);
773-
774774
Counters->ReportCompileRequestCompile(dbCounters);
775775

776776
NWilson::TSpan compileServiceSpan(TWilsonKqp::CompileService, ev->Get() ? std::move(ev->TraceId) : NWilson::TTraceId(), "CompileService");
@@ -785,22 +785,13 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
785785
: (TableServiceConfig.GetEnableAstCache() && !request.QueryAst)
786786
? ECompileActorAction::PARSE
787787
: ECompileActorAction::COMPILE);
788-
auto query = request.Query ? *request.Query : *compileResult->Query;
789-
if (compileResult) {
790-
query.UserSid = compileResult->Query->UserSid;
791-
if (query != *compileResult->Query) {
792-
LOG_WARN_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "queryId in recompile request and queryId in cache are different"
793-
<< ", queryId in request: " << query.SerializeToString()
794-
<< ", queryId in cache: " << compileResult->Query->SerializeToString()
795-
);
796-
}
797-
}
798-
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, query,
788+
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, compileResult ? *compileResult->Query : *request.Query,
799789
compileSettings, request.UserToken, dbCounters, request.GUCSettings, request.ApplicationName,
800790
ev->Cookie, std::move(ev->Get()->IntrestedInResult),
801791
ev->Get()->UserRequestContext,
802792
ev->Get() ? std::move(ev->Get()->Orbit) : NLWTrace::TOrbit(),
803793
std::move(compileServiceSpan), std::move(ev->Get()->TempTablesState));
794+
compileRequest.FindInCache = false;
804795

805796
if (TableServiceConfig.GetEnableAstCache() && request.QueryAst) {
806797
return CompileByAst(*request.QueryAst, compileRequest, ctx);
@@ -981,7 +972,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
981972
<< ", ast: " << queryAst.Ast->Root->ToString());
982973
auto compileResult = QueryCache.FindByAst(compileRequest.Query, *queryAst.Ast, compileRequest.CompileSettings.KeepInCache);
983974

984-
if (HasTempTablesNameClashes(compileResult, compileRequest.TempTablesState)) {
975+
if (!compileRequest.FindInCache || HasTempTablesNameClashes(compileResult, compileRequest.TempTablesState)) {
985976
compileResult = nullptr;
986977
}
987978

@@ -992,7 +983,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
992983
<< ", sender: " << compileRequest.Sender
993984
<< ", queryUid: " << compileResult->Uid);
994985

995-
compileResult->Ast->PgAutoParamValues = std::move(queryAst.Ast->PgAutoParamValues);
986+
compileResult->GetAst()->PgAutoParamValues = std::move(queryAst.Ast->PgAutoParamValues);
996987

997988
ReplyFromCache(compileRequest.Sender, compileResult, ctx, compileRequest.Cookie, std::move(compileRequest.Orbit), std::move(compileRequest.CompileServiceSpan));
998989
return;
@@ -1066,10 +1057,10 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
10661057
if (QueryCache.FindByQuery(query, keepInCache)) {
10671058
return false;
10681059
}
1069-
if (compileResult->Ast && QueryCache.FindByAst(query, *compileResult->Ast, keepInCache)) {
1060+
if (compileResult->GetAst() && QueryCache.FindByAst(query, *compileResult->GetAst(), keepInCache)) {
10701061
return false;
10711062
}
1072-
auto newCompileResult = TKqpCompileResult::Make(CreateGuidAsString(), compileResult->Status, compileResult->Issues, compileResult->MaxReadType, std::move(query), compileResult->Ast);
1063+
auto newCompileResult = TKqpCompileResult::Make(CreateGuidAsString(), compileResult->Status, compileResult->Issues, compileResult->MaxReadType, std::move(query), compileResult->QueryAst);
10731064
newCompileResult->AllowCache = compileResult->AllowCache;
10741065
newCompileResult->PreparedQuery = compileResult->PreparedQuery;
10751066
LOG_DEBUG_S(ctx, NKikimrServices::KQP_COMPILE_SERVICE, "Insert preparing query with params, queryId: " << query.SerializeToString());

‎ydb/core/kqp/session_actor/kqp_query_state.cpp

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ bool TKqpQueryState::SaveAndCheckCompileResult(TEvKqp::TEvCompileResponse* ev) {
161161
CommandTagName = CompileResult->CommandTagName;
162162
}
163163
for (const auto& param : PreparedQuery->GetParameters()) {
164-
const auto& ast = CompileResult->Ast;
164+
const auto& ast = CompileResult->GetAst();
165165
if (!ast || !ast->PgAutoParamValues || !ast->PgAutoParamValues->contains(param.GetName())) {
166166
ResultParams.push_back(param);
167167
}
@@ -276,15 +276,9 @@ std::unique_ptr<TEvKqp::TEvRecompileRequest> TKqpQueryState::BuildReCompileReque
276276
compileDeadline = Min(compileDeadline, QueryDeadlines.CancelAt);
277277
}
278278

279-
TMaybe<TQueryAst> statementAst;
280-
if (!Statements.empty()) {
281-
YQL_ENSURE(CurrentStatementId < Statements.size());
282-
statementAst = Statements[CurrentStatementId];
283-
}
284-
285279
return std::make_unique<TEvKqp::TEvRecompileRequest>(UserToken, CompileResult->Uid, query, isQueryActionPrepare,
286280
compileDeadline, DbCounters, gUCSettingsPtr, ApplicationName, std::move(cookie), UserRequestContext, std::move(Orbit), TempTablesState,
287-
statementAst);
281+
CompileResult->QueryAst);
288282
}
289283

290284
std::unique_ptr<TEvKqp::TEvCompileRequest> TKqpQueryState::BuildSplitRequest(std::shared_ptr<std::atomic<bool>> cookie, const TGUCSettings::TPtr& gUCSettingsPtr) {

‎ydb/core/kqp/session_actor/kqp_session_actor.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -904,8 +904,8 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
904904
try {
905905
const auto& parameters = QueryState->GetYdbParameters();
906906
QueryState->QueryData->ParseParameters(parameters);
907-
if (QueryState->CompileResult && QueryState->CompileResult->Ast && QueryState->CompileResult->Ast->PgAutoParamValues) {
908-
for(const auto& [name, param] : *QueryState->CompileResult->Ast->PgAutoParamValues) {
907+
if (QueryState->CompileResult && QueryState->CompileResult->GetAst() && QueryState->CompileResult->GetAst()->PgAutoParamValues) {
908+
for(const auto& [name, param] : *QueryState->CompileResult->GetAst()->PgAutoParamValues) {
909909
if (!parameters.contains(name)) {
910910
QueryState->QueryData->AddTypedValueParam(name, param);
911911
}

0 commit comments

Comments
 (0)