Skip to content

Commit 0f3b2c1

Browse files
VPolkaVPolka
authored
recompilation with query id from compilation result (#9592)
1 parent de5cc2a commit 0f3b2c1

File tree

6 files changed

+38
-37
lines changed

6 files changed

+38
-37
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 & 14 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
}
@@ -764,8 +766,6 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
764766
}
765767

766768
if (compileResult || request.Query) {
767-
QueryCache.EraseByUid(request.Uid);
768-
769769
Counters->ReportCompileRequestCompile(dbCounters);
770770

771771
NWilson::TSpan compileServiceSpan(TWilsonKqp::CompileService, ev->Get() ? std::move(ev->TraceId) : NWilson::TTraceId(), "CompileService");
@@ -790,12 +790,13 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
790790
);
791791
}
792792
}
793-
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, query,
793+
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, compileResult ? *compileResult->Query : *request.Query,
794794
compileSettings, request.UserToken, dbCounters, request.GUCSettings, request.ApplicationName,
795795
ev->Cookie, std::move(ev->Get()->IntrestedInResult),
796796
ev->Get()->UserRequestContext,
797797
ev->Get() ? std::move(ev->Get()->Orbit) : NLWTrace::TOrbit(),
798798
std::move(compileServiceSpan), std::move(ev->Get()->TempTablesState));
799+
compileRequest.FindInCache = false;
799800

800801
if (TableServiceConfig.GetEnableAstCache() && request.QueryAst) {
801802
return CompileByAst(*request.QueryAst, compileRequest, ctx);
@@ -976,7 +977,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
976977
<< ", ast: " << queryAst.Ast->Root->ToString());
977978
auto compileResult = QueryCache.FindByAst(compileRequest.Query, *queryAst.Ast, compileRequest.CompileSettings.KeepInCache);
978979

979-
if (HasTempTablesNameClashes(compileResult, compileRequest.TempTablesState)) {
980+
if (!compileRequest.FindInCache || HasTempTablesNameClashes(compileResult, compileRequest.TempTablesState)) {
980981
compileResult = nullptr;
981982
}
982983

@@ -987,7 +988,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
987988
<< ", sender: " << compileRequest.Sender
988989
<< ", queryUid: " << compileResult->Uid);
989990

990-
compileResult->Ast->PgAutoParamValues = std::move(queryAst.Ast->PgAutoParamValues);
991+
compileResult->GetAst()->PgAutoParamValues = std::move(queryAst.Ast->PgAutoParamValues);
991992

992993
ReplyFromCache(compileRequest.Sender, compileResult, ctx, compileRequest.Cookie, std::move(compileRequest.Orbit), std::move(compileRequest.CompileServiceSpan));
993994
return;
@@ -1061,10 +1062,10 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
10611062
if (QueryCache.FindByQuery(query, keepInCache)) {
10621063
return false;
10631064
}
1064-
if (compileResult->Ast && QueryCache.FindByAst(query, *compileResult->Ast, keepInCache)) {
1065+
if (compileResult->GetAst() && QueryCache.FindByAst(query, *compileResult->GetAst(), keepInCache)) {
10651066
return false;
10661067
}
1067-
auto newCompileResult = TKqpCompileResult::Make(CreateGuidAsString(), compileResult->Status, compileResult->Issues, compileResult->MaxReadType, std::move(query), compileResult->Ast);
1068+
auto newCompileResult = TKqpCompileResult::Make(CreateGuidAsString(), compileResult->Status, compileResult->Issues, compileResult->MaxReadType, std::move(query), compileResult->QueryAst);
10681069
newCompileResult->AllowCache = compileResult->AllowCache;
10691070
newCompileResult->PreparedQuery = compileResult->PreparedQuery;
10701071
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
@@ -160,7 +160,7 @@ bool TKqpQueryState::SaveAndCheckCompileResult(TEvKqp::TEvCompileResponse* ev) {
160160
CommandTagName = CompileResult->CommandTagName;
161161
}
162162
for (const auto& param : PreparedQuery->GetParameters()) {
163-
const auto& ast = CompileResult->Ast;
163+
const auto& ast = CompileResult->GetAst();
164164
if (!ast || !ast->PgAutoParamValues || !ast->PgAutoParamValues->contains(param.GetName())) {
165165
ResultParams.push_back(param);
166166
}
@@ -275,15 +275,9 @@ std::unique_ptr<TEvKqp::TEvRecompileRequest> TKqpQueryState::BuildReCompileReque
275275
compileDeadline = Min(compileDeadline, QueryDeadlines.CancelAt);
276276
}
277277

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

289283
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
@@ -899,8 +899,8 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
899899
try {
900900
const auto& parameters = QueryState->GetYdbParameters();
901901
QueryState->QueryData->ParseParameters(parameters);
902-
if (QueryState->CompileResult && QueryState->CompileResult->Ast && QueryState->CompileResult->Ast->PgAutoParamValues) {
903-
for(const auto& [name, param] : *QueryState->CompileResult->Ast->PgAutoParamValues) {
902+
if (QueryState->CompileResult && QueryState->CompileResult->GetAst() && QueryState->CompileResult->GetAst()->PgAutoParamValues) {
903+
for(const auto& [name, param] : *QueryState->CompileResult->GetAst()->PgAutoParamValues) {
904904
if (!parameters.contains(name)) {
905905
QueryState->QueryData->AddTypedValueParam(name, param);
906906
}

0 commit comments

Comments
 (0)