Skip to content

Commit bb94eb0

Browse files
gridnevvvitgridnevvvit
committed
parameters pushdown support prototype
1 parent 749b883 commit bb94eb0

File tree

6 files changed

+119
-15
lines changed

6 files changed

+119
-15
lines changed

Diff for: ‎ydb/core/kqp/expr_nodes/kqp_expr_nodes.json

+2-1
Original file line numberDiff line numberDiff line change
@@ -679,7 +679,8 @@
679679
"Children": [
680680
{"Index": 0, "Name": "Type", "Type": "TExprBase"},
681681
{"Index": 1, "Name": "Columns", "Type": "TCoAtomList"},
682-
{"Index": 2, "Name": "Lambda", "Type": "TCoLambda"}
682+
{"Index": 2, "Name": "Parameters", "Type": "TExprList"},
683+
{"Index": 3, "Name": "Lambda", "Type": "TCoLambda"}
683684
]
684685
},
685686
{

Diff for: ‎ydb/core/kqp/host/kqp_type_ann.cpp

+17-2
Original file line numberDiff line numberDiff line change
@@ -1077,7 +1077,7 @@ TStatus AnnotateOlapFilter(const TExprNode::TPtr& node, TExprContext& ctx) {
10771077
}
10781078

10791079
TStatus AnnotateOlapApply(const TExprNode::TPtr& node, TExprContext& ctx) {
1080-
if (!EnsureArgsCount(*node, 3U, ctx)) {
1080+
if (!EnsureArgsCount(*node, 4U, ctx)) {
10811081
return TStatus::Error;
10821082
}
10831083

@@ -1098,6 +1098,7 @@ TStatus AnnotateOlapApply(const TExprNode::TPtr& node, TExprContext& ctx) {
10981098

10991099
const auto structType = argsType->Cast<TStructExprType>();
11001100
TTypeAnnotationNode::TListType argsTypes(columns->ChildrenSize());
1101+
11011102
for (auto i = 0U; i < argsTypes.size(); ++i) {
11021103
if (const auto argType = structType->FindItemType(columns->Child(i)->Content()))
11031104
argsTypes[i] = argType;
@@ -1109,11 +1110,25 @@ TStatus AnnotateOlapApply(const TExprNode::TPtr& node, TExprContext& ctx) {
11091110
}
11101111
}
11111112

1113+
TExprList parameters = TExprList(node->Child(TKqpOlapApply::idx_Parameters));
1114+
std::vector<const NYql::TTypeAnnotationNode *> withParameters = argsTypes;
1115+
for(auto expr: parameters) {
1116+
if (!EnsureArgsCount(*expr.Ptr(), 2U, ctx)) {
1117+
return TStatus::Error;
1118+
}
1119+
1120+
TCoParameter param = TMaybeNode<TCoParameter>(expr.Ptr()).Cast();
1121+
const auto& paramType = expr.Ptr()->Child(TCoParameter::idx_Type);
1122+
if (EnsureType(*paramType, ctx)) {
1123+
withParameters.push_back(paramType->GetTypeAnn()->Cast<TTypeExprType>()->GetType());
1124+
}
1125+
}
1126+
11121127
if (!EnsureLambda(node->Tail(), ctx)) {
11131128
return TStatus::Error;
11141129
}
11151130

1116-
if (!UpdateLambdaAllArgumentsTypes(node->TailRef(), argsTypes, ctx)) {
1131+
if (!UpdateLambdaAllArgumentsTypes(node->TailRef(), withParameters, ctx)) {
11171132
return TStatus::Error;
11181133
}
11191134

Diff for: ‎ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp

+10-8
Original file line numberDiff line numberDiff line change
@@ -214,11 +214,6 @@ TMaybeNode<TExprBase> YqlApplyPushdown(const TExprBase& apply, const TExprNode&
214214
return false;
215215
});
216216

217-
// Temporary fix for https://st.yandex-team.ru/KIKIMR-22216
218-
if (parameters.size()!=0) {
219-
return nullptr;
220-
}
221-
222217
const auto members = FindNodes(apply.Ptr(), [&argument] (const TExprNode::TPtr& node) {
223218
if (const auto maybeMember = TMaybeNode<TCoMember>(node))
224219
return maybeMember.Cast().Struct().Raw() == &argument;
@@ -235,6 +230,12 @@ TMaybeNode<TExprBase> YqlApplyPushdown(const TExprBase& apply, const TExprNode&
235230
replacements.emplace(member.Get(), arguments.back());
236231
}
237232

233+
for(const auto& pptr : parameters) {
234+
TCoParameter parameter = TMaybeNode<TCoParameter>(pptr).Cast();
235+
arguments.emplace_back(ctx.NewArgument(pptr->Pos(), parameter.Name().StringValue()));
236+
replacements.emplace(pptr.Get(), arguments.back());
237+
}
238+
238239
// Temporary fix for https://st.yandex-team.ru/KIKIMR-22560
239240
if (!columns.size()) {
240241
return nullptr;
@@ -243,6 +244,7 @@ TMaybeNode<TExprBase> YqlApplyPushdown(const TExprBase& apply, const TExprNode&
243244
return Build<TKqpOlapApply>(ctx, apply.Pos())
244245
.Type(ExpandType(argument.Pos(), *argument.GetTypeAnn(), ctx))
245246
.Columns().Add(std::move(columns)).Build()
247+
.Parameters().Add(std::move(parameters)).Build()
246248
.Lambda(ctx.NewLambda(apply.Pos(), ctx.NewArguments(argument.Pos(), std::move(arguments)), ctx.ReplaceNodes(apply.Ptr(), replacements)))
247249
.Done();
248250
}
@@ -401,7 +403,7 @@ std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExp
401403
return SafeCastPredicatePushdown(maybeFlatmap.Cast(), argument, ctx, pos);
402404
} else if (auto maybePredicate = node.Maybe<TCoCompare>()) {
403405
return SimplePredicatePushdown(maybePredicate.Cast(), argument, ctx, pos);
404-
}
406+
}
405407

406408
if constexpr (NKikimr::NSsa::RuntimeVersion >= 5U) {
407409
return YqlApplyPushdown(node, argument, ctx);
@@ -800,7 +802,7 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
800802
}
801803
remaining = std::move(remainingAfterApply);
802804
}
803-
805+
804806
if (pushedPredicates.empty()) {
805807
return node;
806808
}
@@ -810,7 +812,7 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
810812
const auto remainingFilter = CombinePredicatesWithAnd(remaining, ctx, node.Pos(), false, true);
811813

812814
TMaybeNode<TExprBase> olapFilter;
813-
if (pushedFilter.FirstLevelOps.IsValid()) {
815+
if (pushedFilter.FirstLevelOps.IsValid()) {
814816
olapFilter = Build<TKqpOlapFilter>(ctx, node.Pos())
815817
.Input(read.Process().Body())
816818
.Condition(pushedFilter.FirstLevelOps.Cast())

Diff for: ‎ydb/core/kqp/query_compiler/kqp_olap_compiler.cpp

+6
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,12 @@ TTypedColumn CompileYqlKernelScalarApply(const TKqpOlapApply& apply, TKqpOlapCom
600600
argTypes.emplace_back(arg.Type);
601601
}
602602

603+
for(const auto& param: apply.Parameters()) {
604+
const auto& arg = GetOrCreateColumnIdAndType(param, ctx);
605+
ids.emplace_back(arg.Id);
606+
argTypes.emplace_back(arg.Type);
607+
}
608+
603609
auto *const command = ctx.CreateAssignCmd();
604610
auto *const function = command->MutableFunction();
605611
const auto idx = ctx.GetKernelRequestBuilder().AddScalarApply(apply.Lambda().Ref(), argTypes, ctx.ExprCtx());

Diff for: ‎ydb/core/kqp/ut/olap/kqp_olap_ut.cpp

+83-4
Original file line numberDiff line numberDiff line change
@@ -1458,7 +1458,7 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
14581458
//This check is based on an assumpltion, that for pushed down predicates column names are preserved in AST
14591459
//But for non-pushed down predicates column names are (usually) replaced with a label, started with $. It' not a rule, but a heuristic
14601460
//So this check may require a correction when some ast optimization rules are changed
1461-
UNIT_ASSERT_C(ast.find(R"((Unwrap (/ $)") != std::string::npos,
1461+
UNIT_ASSERT_C(ast.find(R"((Unwrap (/ $)") != std::string::npos,
14621462
TStringBuilder() << "Unsafe subpredicate is pushed down. Query: " << query);
14631463

14641464
UNIT_ASSERT_C(ast.find("NarrowMap") != std::string::npos,
@@ -1776,6 +1776,85 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
17761776
UNIT_ASSERT_VALUES_EQUAL(result, insertRows);
17771777
}
17781778

1779+
Y_UNIT_TEST(PredicatePushdownWithParametersILike) {
1780+
constexpr bool logQueries = true;
1781+
auto settings = TKikimrSettings()
1782+
.SetWithSampleTables(false);
1783+
TKikimrRunner kikimr(settings);
1784+
1785+
TStreamExecScanQuerySettings scanSettings;
1786+
scanSettings.Explain(true);
1787+
1788+
TLocalHelper(kikimr.GetTestServer()).CreateTestOlapTable();
1789+
WriteTestData(kikimr, "/Root/olapStore/olapTable", 10000, 3000000, 1000);
1790+
1791+
auto tableClient = kikimr.GetTableClient();
1792+
1793+
auto buildQuery = [](bool pushEnabled) {
1794+
TStringBuilder builder;
1795+
1796+
builder << "--!syntax_v1" << Endl;
1797+
1798+
if (!pushEnabled) {
1799+
builder << "PRAGMA Kikimr.OptEnableOlapPushdown = \"false\";" << Endl;
1800+
}
1801+
1802+
builder << R"(
1803+
DECLARE $in_uid AS Utf8;
1804+
DECLARE $in_level AS Int32;
1805+
1806+
SELECT `timestamp` FROM `/Root/olapStore/olapTable` WHERE
1807+
uid ILIKE "uid_" || $in_uid || "980" AND level > $in_level
1808+
ORDER BY `timestamp`;
1809+
)" << Endl;
1810+
1811+
return builder;
1812+
};
1813+
1814+
auto normalQuery = buildQuery(false);
1815+
auto pushQuery = buildQuery(true);
1816+
1817+
auto params = tableClient.GetParamsBuilder()
1818+
.AddParam("$in_uid")
1819+
.Utf8("3000")
1820+
.Build()
1821+
.AddParam("$in_level")
1822+
.Int32(2)
1823+
.Build()
1824+
.Build();
1825+
1826+
auto it = tableClient.StreamExecuteScanQuery(normalQuery, params).GetValueSync();
1827+
UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
1828+
auto goodResult = CollectStreamResult(it);
1829+
1830+
it = tableClient.StreamExecuteScanQuery(pushQuery, params).GetValueSync();
1831+
UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
1832+
auto pushResult = CollectStreamResult(it);
1833+
1834+
if (logQueries) {
1835+
Cerr << "Query: " << normalQuery << Endl;
1836+
Cerr << "Expected: " << goodResult.ResultSetYson << Endl;
1837+
Cerr << "Received: " << pushResult.ResultSetYson << Endl;
1838+
}
1839+
1840+
CompareYson(goodResult.ResultSetYson, pushResult.ResultSetYson);
1841+
1842+
it = tableClient.StreamExecuteScanQuery(pushQuery, scanSettings).GetValueSync();
1843+
UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
1844+
1845+
auto result = CollectStreamResult(it);
1846+
auto ast = result.QueryStats->Getquery_ast();
1847+
1848+
UNIT_ASSERT_C(ast.find("KqpOlapFilter") != std::string::npos,
1849+
TStringBuilder() << "Predicate not pushed down. Query: " << pushQuery);
1850+
1851+
NJson::TJsonValue plan, readRange;
1852+
NJson::ReadJsonTree(*result.PlanJson, &plan, true);
1853+
1854+
readRange = FindPlanNodeByKv(plan, "Name", "TableFullScan");
1855+
UNIT_ASSERT(readRange.IsDefined());
1856+
}
1857+
17791858
Y_UNIT_TEST(PredicatePushdownWithParameters) {
17801859
constexpr bool logQueries = true;
17811860
auto settings = TKikimrSettings()
@@ -2052,7 +2131,7 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
20522131
Age Int16 NOT NULL,
20532132
Sex Int16 NOT NULL,
20542133
PRIMARY KEY (CounterID, WatchID)
2055-
)
2134+
)а
20562135
PARTITION BY HASH(WatchID)
20572136
WITH (
20582137
STORE = COLUMN,
@@ -3277,9 +3356,9 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
32773356
32783357
)", noTx).GetValueSync();
32793358
UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
3280-
3359+
32813360
result = queryClient.ExecuteQuery(R"(
3282-
UPSERT INTO Test (Id, Name, Comment) VALUES
3361+
UPSERT INTO Test (Id, Name, Comment) VALUES
32833362
(10, "n1", "aa"),
32843363
(20, "n2", "bb"),
32853364
(30, "n3", "cc"),

Diff for: ‎yql/essentials/core/yql_expr_type_annotation.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -2025,6 +2025,7 @@ bool IsDataOrOptionalOfDataOrPg(const TTypeAnnotationNode* typeAnnotation) {
20252025

20262026
bool EnsureArgsCount(const TExprNode& node, ui32 expectedArgs, TExprContext& ctx) {
20272027
if (node.ChildrenSize() != expectedArgs) {
2028+
Y_ABORT_UNLESS(false, "text");
20282029
ctx.AddError(TIssue(ctx.GetPosition(node.Pos()), TStringBuilder() << "Expected " << expectedArgs << " argument(s), but got " <<
20292030
node.ChildrenSize()));
20302031
return false;

0 commit comments

Comments
 (0)