Skip to content

Commit 903def9

Browse files
pavelvelikhovpavelvelikhov
authored
Generate single JSON with all types of plans in the server (#1606)
* Moved plan simplification to the server * Small fix in plan generation * Minor fix in scripting * Don't print OLAP plans yet * Fixed unused arg problem * fixed a typo * Simplifying a copy of the plan now * Fixed minor bugs * fixed a bug in ut_common * Cannonized some tests, changed plan comparison
1 parent b71a6fd commit 903def9

File tree

14 files changed

+1142
-1116
lines changed

14 files changed

+1142
-1116
lines changed

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

Lines changed: 220 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
#include <ydb/library/yql/utils/plan/plan_utils.h>
1717
#include <ydb/library/yql/providers/dq/common/yql_dq_settings.h>
1818

19+
#include <ydb/public/lib/ydb_cli/common/format.h>
20+
1921
#include <library/cpp/json/writer/json.h>
2022
#include <library/cpp/json/json_reader.h>
2123
#include <library/cpp/protobuf/json/proto2json.h>
@@ -1811,6 +1813,217 @@ void SetNonZero(NJson::TJsonValue& node, const TStringBuf& name, T value) {
18111813
}
18121814
}
18131815

1816+
void BuildPlanIndex(NJson::TJsonValue& plan, THashMap<int, NJson::TJsonValue>& planIndex, THashMap<TString, NJson::TJsonValue>& precomputes) {
1817+
if (plan.GetMapSafe().contains("PlanNodeId")){
1818+
auto id = plan.GetMapSafe().at("PlanNodeId").GetIntegerSafe();
1819+
planIndex[id] = plan;
1820+
}
1821+
1822+
if (plan.GetMapSafe().contains("Subplan Name")) {
1823+
const auto& precomputeName = plan.GetMapSafe().at("Subplan Name").GetStringSafe();
1824+
1825+
auto pos = precomputeName.find("precompute");
1826+
if (pos != TString::npos) {
1827+
precomputes[precomputeName.substr(pos)] = plan;
1828+
}
1829+
}
1830+
1831+
if (plan.GetMapSafe().contains("Plans")) {
1832+
for (auto p : plan.GetMapSafe().at("Plans").GetArraySafe()) {
1833+
BuildPlanIndex(p, planIndex, precomputes);
1834+
}
1835+
}
1836+
}
1837+
1838+
TVector<NJson::TJsonValue> RemoveRedundantNodes(NJson::TJsonValue& plan, const THashSet<TString>& redundantNodes) {
1839+
auto& planMap = plan.GetMapSafe();
1840+
1841+
TVector<NJson::TJsonValue> children;
1842+
if (planMap.contains("Plans") && planMap.at("Plans").IsArray()) {
1843+
for (auto& child : planMap.at("Plans").GetArraySafe()) {
1844+
auto newChildren = RemoveRedundantNodes(child, redundantNodes);
1845+
children.insert(children.end(), newChildren.begin(), newChildren.end());
1846+
}
1847+
}
1848+
1849+
planMap.erase("Plans");
1850+
if (!children.empty()) {
1851+
auto& plans = planMap["Plans"];
1852+
for (auto& child : children) {
1853+
plans.AppendValue(child);
1854+
}
1855+
}
1856+
1857+
const auto typeName = planMap.at("Node Type").GetStringSafe();
1858+
if (redundantNodes.contains(typeName) || typeName.find("Precompute") != TString::npos) {
1859+
return children;
1860+
}
1861+
1862+
return {plan};
1863+
}
1864+
1865+
NJson::TJsonValue ReconstructQueryPlanRec(const NJson::TJsonValue& plan,
1866+
int operatorIndex,
1867+
const THashMap<int, NJson::TJsonValue>& planIndex,
1868+
const THashMap<TString, NJson::TJsonValue>& precomputes,
1869+
int& nodeCounter) {
1870+
1871+
int currentNodeId = nodeCounter++;
1872+
1873+
NJson::TJsonValue result;
1874+
result["PlanNodeId"] = currentNodeId;
1875+
if (plan.GetMapSafe().contains("PlanNodeType")) {
1876+
result["PlanNodeType"] = plan.GetMapSafe().at("PlanNodeType").GetStringSafe();
1877+
}
1878+
1879+
if (plan.GetMapSafe().contains("Stats")) {
1880+
result["Stats"] = plan.GetMapSafe().at("Stats");
1881+
}
1882+
1883+
if (!plan.GetMapSafe().contains("Operators")) {
1884+
NJson::TJsonValue planInputs;
1885+
1886+
result["Node Type"] = plan.GetMapSafe().at("Node Type").GetStringSafe();
1887+
1888+
if (!plan.GetMapSafe().contains("Plans")) {
1889+
return result;
1890+
}
1891+
1892+
if (plan.GetMapSafe().at("Node Type") == "TableLookup") {
1893+
NJson::TJsonValue newOps;
1894+
NJson::TJsonValue op;
1895+
1896+
op["Name"] = "TableLookup";
1897+
op["Columns"] = plan.GetMapSafe().at("Columns");
1898+
op["LookupKeyColumns"] = plan.GetMapSafe().at("LookupKeyColumns");
1899+
op["Table"] = plan.GetMapSafe().at("Table");
1900+
1901+
newOps.AppendValue(op);
1902+
1903+
result["Operators"] = newOps;
1904+
return result;
1905+
}
1906+
1907+
for (auto p : plan.GetMapSafe().at("Plans").GetArraySafe()) {
1908+
if (p.GetMapSafe().at("Node Type").GetStringSafe().find("Precompute") == TString::npos) {
1909+
planInputs.AppendValue(ReconstructQueryPlanRec(p, 0, planIndex, precomputes, nodeCounter));
1910+
}
1911+
}
1912+
result["Plans"] = planInputs;
1913+
return result;
1914+
}
1915+
1916+
if (plan.GetMapSafe().contains("CTE Name") && plan.GetMapSafe().at("Node Type") == "ConstantExpr") {
1917+
auto precompute = plan.GetMapSafe().at("CTE Name").GetStringSafe();
1918+
if (!precomputes.contains(precompute)) {
1919+
result["Node Type"] = "ConstantExpr";
1920+
return result;
1921+
}
1922+
return ReconstructQueryPlanRec(precomputes.at(precompute), 0, planIndex, precomputes, nodeCounter);
1923+
}
1924+
1925+
auto ops = plan.GetMapSafe().at("Operators").GetArraySafe();
1926+
auto op = ops[operatorIndex];
1927+
1928+
TVector<NJson::TJsonValue> planInputs;
1929+
1930+
auto opName = op.GetMapSafe().at("Name").GetStringSafe();
1931+
1932+
for (auto opInput : op.GetMapSafe().at("Inputs").GetArraySafe()) {
1933+
if (opInput.GetMapSafe().contains("ExternalPlanNodeId")) {
1934+
auto inputPlanKey = opInput.GetMapSafe().at("ExternalPlanNodeId").GetIntegerSafe();
1935+
auto inputPlan = planIndex.at(inputPlanKey);
1936+
planInputs.push_back( ReconstructQueryPlanRec(inputPlan, 0, planIndex, precomputes, nodeCounter));
1937+
} else if (opInput.GetMapSafe().contains("InternalOperatorId")) {
1938+
auto inputPlanId = opInput.GetMapSafe().at("InternalOperatorId").GetIntegerSafe();
1939+
planInputs.push_back( ReconstructQueryPlanRec(plan, inputPlanId, planIndex, precomputes, nodeCounter));
1940+
}
1941+
// temp hack
1942+
if (opName == "Filter") {
1943+
break;
1944+
}
1945+
}
1946+
1947+
if (op.GetMapSafe().contains("Inputs")) {
1948+
op.GetMapSafe().erase("Inputs");
1949+
}
1950+
1951+
if (op.GetMapSafe().contains("Input") || op.GetMapSafe().contains("ToFlow")) {
1952+
TString maybePrecompute = "";
1953+
if (op.GetMapSafe().contains("Input")) {
1954+
maybePrecompute = op.GetMapSafe().at("Input").GetStringSafe();
1955+
} else if (op.GetMapSafe().contains("ToFlow")) {
1956+
maybePrecompute = op.GetMapSafe().at("ToFlow").GetStringSafe();
1957+
}
1958+
1959+
if (precomputes.contains(maybePrecompute)) {
1960+
planInputs.push_back(ReconstructQueryPlanRec(precomputes.at(maybePrecompute), 0, planIndex, precomputes, nodeCounter));
1961+
}
1962+
}
1963+
1964+
result["Node Type"] = opName;
1965+
NJson::TJsonValue newOps;
1966+
newOps.AppendValue(op);
1967+
result["Operators"] = newOps;
1968+
1969+
if (planInputs.size()){
1970+
NJson::TJsonValue plans;
1971+
for( auto i : planInputs) {
1972+
plans.AppendValue(i);
1973+
}
1974+
result["Plans"] = plans;
1975+
}
1976+
1977+
return result;
1978+
}
1979+
1980+
NJson::TJsonValue SimplifyQueryPlan(NJson::TJsonValue& plan) {
1981+
static const THashSet<TString> redundantNodes = {
1982+
"UnionAll",
1983+
"Broadcast",
1984+
"Map",
1985+
"HashShuffle",
1986+
"Merge",
1987+
"Collect",
1988+
"Stage",
1989+
"Iterator",
1990+
"PartitionByKey",
1991+
"ToFlow"
1992+
};
1993+
1994+
THashMap<int, NJson::TJsonValue> planIndex;
1995+
THashMap<TString, NJson::TJsonValue> precomputes;
1996+
1997+
1998+
BuildPlanIndex(plan, planIndex, precomputes);
1999+
2000+
int nodeCounter = 0;
2001+
plan = ReconstructQueryPlanRec(plan, 0, planIndex, precomputes, nodeCounter);
2002+
RemoveRedundantNodes(plan, redundantNodes);
2003+
return plan;
2004+
}
2005+
2006+
TString AddSimplifiedPlan(const TString& planText, bool analyzeMode) {
2007+
Y_UNUSED(analyzeMode);
2008+
NJson::TJsonValue planJson;
2009+
NJson::ReadJsonTree(planText, &planJson, true);
2010+
if (!planJson.GetMapSafe().contains("Plan")){
2011+
return planText;
2012+
}
2013+
2014+
NJson::TJsonValue planCopy;
2015+
NJson::ReadJsonTree(planText, &planCopy, true);
2016+
2017+
planJson["SimplifiedPlan"] = SimplifyQueryPlan(planCopy.GetMapSafe().at("Plan"));
2018+
2019+
// Don't print the OLAP plan yet, there are some non UTF-8 symbols there that need to be fixed
2020+
//TTempBufOutput stringStream;
2021+
//NYdb::NConsoleClient::TQueryPlanPrinter printer(NYdb::NConsoleClient::EOutputFormat::PrettyTable, analyzeMode, stringStream);
2022+
//printer.Print(planJson.GetStringRobust());
2023+
//planJson["OLAPText"] = stringStream.Data();
2024+
return planJson.GetStringRobust();
2025+
}
2026+
18142027
TString SerializeTxPlans(const TVector<const TString>& txPlans, const TString commonPlanInfo = "") {
18152028
NJsonWriter::TBuf writer;
18162029
writer.SetIndentSpaces(2);
@@ -1862,7 +2075,8 @@ TString SerializeTxPlans(const TVector<const TString>& txPlans, const TString co
18622075
writer.EndObject();
18632076
writer.EndObject();
18642077

1865-
return writer.Str();
2078+
auto resultPlan = writer.Str();
2079+
return AddSimplifiedPlan(resultPlan, false);
18662080
}
18672081

18682082
} // namespace
@@ -2250,7 +2464,8 @@ TString AddExecStatsToTxPlan(const TString& txPlanJson, const NYql::NDqProto::TD
22502464

22512465
NJsonWriter::TBuf txWriter;
22522466
txWriter.WriteJsonValue(&root, true);
2253-
return txWriter.Str();
2467+
auto resultPlan = txWriter.Str();
2468+
return AddSimplifiedPlan(resultPlan, true);
22542469
}
22552470

22562471
TString SerializeAnalyzePlan(const NKqpProto::TKqpStatsQuery& queryStats) {
@@ -2294,6 +2509,9 @@ TString SerializeScriptPlan(const TVector<const TString>& queryPlans) {
22942509
if (auto dqPlan = planMap.FindPtr("Plan")) {
22952510
writer.WriteKey("Plan");
22962511
writer.WriteJsonValue(dqPlan);
2512+
writer.WriteKey("SimplifiedPlan");
2513+
auto simplifiedPlan = SimplifyQueryPlan(*dqPlan);
2514+
writer.WriteJsonValue(&simplifiedPlan);
22972515
}
22982516
writer.EndObject();
22992517
}

‎ydb/core/kqp/ut/common/kqp_ut_common.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1170,7 +1170,7 @@ std::vector<NJson::TJsonValue> FindPlanNodes(const NJson::TJsonValue& plan, cons
11701170

11711171
std::vector<NJson::TJsonValue> FindPlanStages(const NJson::TJsonValue& plan) {
11721172
std::vector<NJson::TJsonValue> stages;
1173-
FindPlanStagesImpl(plan, stages);
1173+
FindPlanStagesImpl(plan.GetMapSafe().at("Plan"), stages);
11741174
return stages;
11751175
}
11761176

0 commit comments

Comments
 (0)