Skip to content

Commit a26ebd7

Browse files
pavelvelikhovpavelvelikhov
authored
Fixed a problem with simplified plan JSONs (#1878)
1 parent af4bda3 commit a26ebd7

1 file changed

Lines changed: 107 additions & 9 deletions

File tree

‎ydb/core/kqp/opt/kqp_query_plan.cpp‎

Lines changed: 107 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
#include <unordered_map>
3030
#include <regex>
3131

32+
#include <ydb/library/yql/utils/log/log.h>
33+
34+
3235
namespace NKikimr {
3336
namespace NKqp {
3437

@@ -1015,6 +1018,10 @@ class TxPlanSerializer {
10151018
operatorId = Visit(maybeExtend.Cast(), planNode);
10161019
} else if (auto maybeToFlow = TMaybeNode<TCoToFlow>(node)) {
10171020
operatorId = Visit(maybeToFlow.Cast(), planNode);
1021+
} else if (auto maybeAssumeSorted = TMaybeNode<TCoAssumeSorted>(node)) {
1022+
operatorId = Visit(maybeAssumeSorted.Cast(), planNode);
1023+
} else if (auto maybeMember = TMaybeNode<TCoMember>(node)) {
1024+
operatorId = Visit(maybeMember.Cast(), planNode);
10181025
} else if (auto maybeIter = TMaybeNode<TCoIterator>(node)) {
10191026
operatorId = Visit(maybeIter.Cast(), planNode);
10201027
} else if (auto maybePartitionByKey = TMaybeNode<TCoPartitionByKey>(node)) {
@@ -1183,6 +1190,56 @@ class TxPlanSerializer {
11831190
return AddOperator(planNode, "ConstantExpr", std::move(op));
11841191
}
11851192

1193+
TMaybe<std::variant<ui32, TArgContext>> Visit(const TCoAssumeSorted& assumeSorted, TQueryPlanNode& planNode) {
1194+
const auto precomputeValue = NPlanUtils::PrettyExprStr(assumeSorted.Input());
1195+
1196+
TOperator op;
1197+
op.Properties["Name"] = "AssumeSorted";
1198+
1199+
if (auto maybeResultBinding = ContainResultBinding(precomputeValue)) {
1200+
auto [txId, resId] = *maybeResultBinding;
1201+
planNode.CteRefName = TStringBuilder() << "precompute_" << txId << "_" << resId;
1202+
op.Properties["AssumeSorted"] = *planNode.CteRefName;
1203+
} else {
1204+
auto inputs = Visit(assumeSorted.Input().Ptr(), planNode);
1205+
if (inputs.size() == 1) {
1206+
return inputs[0];
1207+
} else {
1208+
return TMaybe<std::variant<ui32, TArgContext>> ();
1209+
}
1210+
}
1211+
1212+
return AddOperator(planNode, "ConstantExpr", std::move(op));
1213+
}
1214+
1215+
TMaybe<std::variant<ui32, TArgContext>> Visit(const TCoMember& member, TQueryPlanNode& planNode) {
1216+
const auto memberValue = NPlanUtils::PrettyExprStr(member.Struct());
1217+
1218+
TOperator op;
1219+
op.Properties["Name"] = "Member";
1220+
1221+
if (auto maybeResultBinding = ContainResultBinding(memberValue)) {
1222+
auto [txId, resId] = *maybeResultBinding;
1223+
planNode.CteRefName = TStringBuilder() << "precompute_" << txId << "_" << resId;
1224+
op.Properties["Member"] = *planNode.CteRefName;
1225+
} else {
1226+
auto inputs = Visit(member.Struct().Ptr(), planNode);
1227+
if (inputs.size() == 1) {
1228+
return inputs[0];
1229+
} else {
1230+
return TMaybe<std::variant<ui32, TArgContext>> ();
1231+
}
1232+
}
1233+
1234+
if (std::find_if(planNode.Operators.begin(), planNode.Operators.end(), [op](const auto & x) {
1235+
return x.Properties.at("Name")=="Member" && x.Properties.at("Member")==op.Properties.at("Member");
1236+
})) {
1237+
return TMaybe<std::variant<ui32, TArgContext>> ();
1238+
}
1239+
1240+
return AddOperator(planNode, "ConstantExpr", std::move(op));
1241+
}
1242+
11861243
std::variant<ui32, TArgContext> Visit(const TCoIterator& iter, TQueryPlanNode& planNode) {
11871244
const auto iterValue = NPlanUtils::PrettyExprStr(iter.List());
11881245

@@ -1825,6 +1882,8 @@ void BuildPlanIndex(NJson::TJsonValue& plan, THashMap<int, NJson::TJsonValue>& p
18251882
auto pos = precomputeName.find("precompute");
18261883
if (pos != TString::npos) {
18271884
precomputes[precomputeName.substr(pos)] = plan;
1885+
} else if (precomputeName.size()>=4 && precomputeName.find("CTE ") != TString::npos) {
1886+
precomputes[precomputeName.substr(4)] = plan;
18281887
}
18291888
}
18301889

@@ -1872,6 +1931,11 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
18721931

18731932
NJson::TJsonValue result;
18741933
result["PlanNodeId"] = currentNodeId;
1934+
1935+
//if (plan.GetMapSafe().contains("PlanNodeId")) {
1936+
// YQL_CLOG(TRACE, CoreDq) << "Recursed into " << plan.GetMapSafe().at("PlanNodeId").GetIntegerSafe() << ", constructed: " << currentNodeId;
1937+
//}
1938+
18751939
if (plan.GetMapSafe().contains("PlanNodeType")) {
18761940
result["PlanNodeType"] = plan.GetMapSafe().at("PlanNodeType").GetStringSafe();
18771941
}
@@ -1889,7 +1953,7 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
18891953
return result;
18901954
}
18911955

1892-
if (plan.GetMapSafe().at("Node Type") == "TableLookup") {
1956+
if (plan.GetMapSafe().at("Node Type").GetStringSafe() == "TableLookup") {
18931957
NJson::TJsonValue newOps;
18941958
NJson::TJsonValue op;
18951959

@@ -1905,20 +1969,32 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
19051969
}
19061970

19071971
for (auto p : plan.GetMapSafe().at("Plans").GetArraySafe()) {
1908-
if (p.GetMapSafe().at("Node Type").GetStringSafe().find("Precompute") == TString::npos) {
1972+
if (!p.GetMapSafe().contains("Operators") && p.GetMapSafe().contains("CTE Name")) {
1973+
auto precompute = p.GetMapSafe().at("CTE Name").GetStringSafe();
1974+
if (precomputes.contains(precompute)) {
1975+
//YQL_CLOG(TRACE, CoreDq) << "Following precompute: " << precompute ;
1976+
planInputs.AppendValue(ReconstructQueryPlanRec(precomputes.at(precompute), 0, planIndex, precomputes, nodeCounter));
1977+
} //else {
1978+
// YQL_CLOG(TRACE, CoreDq) << "Didn't find precompute: " << precompute ;
1979+
//}
1980+
} else if (p.GetMapSafe().at("Node Type").GetStringSafe().find("Precompute") == TString::npos) {
19091981
planInputs.AppendValue(ReconstructQueryPlanRec(p, 0, planIndex, precomputes, nodeCounter));
19101982
}
19111983
}
19121984
result["Plans"] = planInputs;
19131985
return result;
19141986
}
19151987

1916-
if (plan.GetMapSafe().contains("CTE Name") && plan.GetMapSafe().at("Node Type") == "ConstantExpr") {
1988+
if (plan.GetMapSafe().contains("CTE Name") && plan.GetMapSafe().at("Node Type").GetStringSafe() == "ConstantExpr") {
19171989
auto precompute = plan.GetMapSafe().at("CTE Name").GetStringSafe();
19181990
if (!precomputes.contains(precompute)) {
1919-
result["Node Type"] = "ConstantExpr";
1991+
//YQL_CLOG(TRACE, CoreDq) << "Didn't find precompute: " << precompute ;
1992+
1993+
result["Node Type"] = plan.GetMapSafe().at("Node Type");
19201994
return result;
19211995
}
1996+
//YQL_CLOG(TRACE, CoreDq) << "Following precompute: " << precompute ;
1997+
19221998
return ReconstructQueryPlanRec(precomputes.at(precompute), 0, planIndex, precomputes, nodeCounter);
19231999
}
19242000

@@ -1930,6 +2006,11 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
19302006
auto opName = op.GetMapSafe().at("Name").GetStringSafe();
19312007

19322008
for (auto opInput : op.GetMapSafe().at("Inputs").GetArraySafe()) {
2009+
// Sometimes we have inputs for these operators, don't process them
2010+
if (opName == "TablePointLookup") {
2011+
break;
2012+
}
2013+
19332014
if (opInput.GetMapSafe().contains("ExternalPlanNodeId")) {
19342015
auto inputPlanKey = opInput.GetMapSafe().at("ExternalPlanNodeId").GetIntegerSafe();
19352016
auto inputPlan = planIndex.at(inputPlanKey);
@@ -1938,8 +2019,9 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
19382019
auto inputPlanId = opInput.GetMapSafe().at("InternalOperatorId").GetIntegerSafe();
19392020
planInputs.push_back( ReconstructQueryPlanRec(plan, inputPlanId, planIndex, precomputes, nodeCounter));
19402021
}
1941-
// temp hack
1942-
if (opName == "Filter") {
2022+
2023+
// Sometimes we have multiple inputs for these operators, break after the first one
2024+
if (opName == "Filter" || opName == "TopSort") {
19432025
break;
19442026
}
19452027
}
@@ -1948,17 +2030,31 @@ NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
19482030
op.GetMapSafe().erase("Inputs");
19492031
}
19502032

1951-
if (op.GetMapSafe().contains("Input") || op.GetMapSafe().contains("ToFlow")) {
2033+
if (op.GetMapSafe().contains("Input")
2034+
|| op.GetMapSafe().contains("ToFlow")
2035+
|| op.GetMapSafe().contains("Member")
2036+
|| op.GetMapSafe().contains("AssumeSorted")
2037+
|| op.GetMapSafe().contains("Iterator")) {
2038+
19522039
TString maybePrecompute = "";
19532040
if (op.GetMapSafe().contains("Input")) {
19542041
maybePrecompute = op.GetMapSafe().at("Input").GetStringSafe();
19552042
} else if (op.GetMapSafe().contains("ToFlow")) {
19562043
maybePrecompute = op.GetMapSafe().at("ToFlow").GetStringSafe();
2044+
} else if (op.GetMapSafe().contains("Member")) {
2045+
maybePrecompute = op.GetMapSafe().at("Member").GetStringSafe();
2046+
} else if (op.GetMapSafe().contains("AssumeSorted")) {
2047+
maybePrecompute = op.GetMapSafe().at("AssumeSorted").GetStringSafe();
2048+
} else if (op.GetMapSafe().contains("Iterator")) {
2049+
maybePrecompute = op.GetMapSafe().at("Iterator").GetStringSafe();
19572050
}
19582051

19592052
if (precomputes.contains(maybePrecompute)) {
2053+
//YQL_CLOG(TRACE, CoreDq) << "Following precompute: " << maybePrecompute ;
19602054
planInputs.push_back(ReconstructQueryPlanRec(precomputes.at(maybePrecompute), 0, planIndex, precomputes, nodeCounter));
1961-
}
2055+
} //else {
2056+
// YQL_CLOG(TRACE, CoreDq) << "Didn't find precompute: " << maybePrecompute ;
2057+
//}
19622058
}
19632059

19642060
result["Node Type"] = opName;
@@ -1988,7 +2084,9 @@ NJson::TJsonValue SimplifyQueryPlan(NJson::TJsonValue& plan) {
19882084
"Stage",
19892085
"Iterator",
19902086
"PartitionByKey",
1991-
"ToFlow"
2087+
"ToFlow",
2088+
"Member",
2089+
"AssumeSorted"
19922090
};
19932091

19942092
THashMap<int, NJson::TJsonValue> planIndex;

0 commit comments

Comments
 (0)