Skip to content

Commit 6b690bd

Browse files
vitstnvitstn
authored
Merge 879778c into 487c4eb 
2 parents 487c4eb + 879778c commit 6b690bd

File tree

12 files changed

+255
-0
lines changed

12 files changed

+255
-0
lines changed

‎ydb/library/yql/core/common_opt/yql_co_pgselect.cpp‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4154,4 +4154,123 @@ TExprNode::TPtr ExpandPgGrouping(const TExprNode::TPtr& node, TExprContext& ctx,
41544154
.Build();
41554155
}
41564156

4157+
TExprNode::TPtr ExpandPgIterate(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
4158+
const bool all = node->Content().EndsWith("All");
4159+
auto init = node->HeadPtr();
4160+
init = ctx.WrapByCallableIf(!all, "ListUniq", std::move(init));
4161+
auto lambda = node->TailPtr();
4162+
const auto limit = optCtx.Types->PgIterateLimit;
4163+
auto itemArg = ctx.NewArgument(node->Pos(), "item");
4164+
auto stateArg = ctx.NewArgument(node->Pos(), "state");
4165+
auto state = ctx.Builder(node->Pos())
4166+
.Callable("Ensure")
4167+
.Add(0, stateArg)
4168+
.Callable(1, "<")
4169+
.Add(0, itemArg)
4170+
.Callable(1, "Uint32")
4171+
.Atom(0, limit)
4172+
.Seal()
4173+
.Seal()
4174+
.Callable(2, "String")
4175+
.Atom(0, "Too many CTE iterations: " + ToString(limit))
4176+
.Seal()
4177+
.Seal()
4178+
.Build();
4179+
4180+
auto currentIter = ctx.NewCallable(node->Pos(), "Nth", { state, ctx.NewAtom(node->Pos(), 1)});
4181+
auto add = ctx.Builder(node->Pos())
4182+
.Apply(lambda)
4183+
.With(0, currentIter)
4184+
.Seal()
4185+
.Build();
4186+
4187+
auto currentRes = ctx.NewCallable(node->Pos(), "Nth", { state, ctx.NewAtom(node->Pos(), 0)});
4188+
add = ctx.WrapByCallableIf(!all, "ListUniq", std::move(add));
4189+
if (!all) {
4190+
auto id = ctx.Builder(node->Pos())
4191+
.Lambda()
4192+
.Param("x")
4193+
.Arg("x")
4194+
.Seal()
4195+
.Build();
4196+
4197+
add = ctx.Builder(node->Pos())
4198+
.Callable("Filter")
4199+
.Add(0, add)
4200+
.Lambda(1)
4201+
.Param("item")
4202+
.Callable("Not")
4203+
.Callable(0, "Contains")
4204+
.Callable(0, "ToDict")
4205+
.Add(0, currentRes)
4206+
.Add(1, id)
4207+
.Add(2, id)
4208+
.List(3)
4209+
.Atom(0, "Auto", TNodeFlags::Default)
4210+
.Atom(1, "One", TNodeFlags::Default)
4211+
.Seal()
4212+
.Seal()
4213+
.Arg(1, "item")
4214+
.Seal()
4215+
.Seal()
4216+
.Seal()
4217+
.Seal()
4218+
.Build();
4219+
}
4220+
4221+
auto res = ctx.NewCallable(node->Pos(), "Extend", { currentRes, add });
4222+
auto p = ctx.NewList(node->Pos(), { res, add });
4223+
auto foldLambdaBody = ctx.NewList(node->Pos(), { p, p});
4224+
auto foldLambda = ctx.NewLambda(node->Pos(), ctx.NewArguments(node->Pos(), { itemArg, stateArg}), std::move(foldLambdaBody));
4225+
auto foldMap = ctx.Builder(node->Pos())
4226+
.Callable("FoldMap")
4227+
.Callable(0, "ListFromRange")
4228+
.Callable(0, "Uint32")
4229+
.Atom(0, 0)
4230+
.Seal()
4231+
.Callable(1, "Uint32")
4232+
.Atom(0, limit + 1)
4233+
.Seal()
4234+
.Seal()
4235+
.List(1)
4236+
.Add(0, init)
4237+
.Add(1, init)
4238+
.Seal()
4239+
.Add(2, foldLambda)
4240+
.Seal()
4241+
.Build();
4242+
4243+
return ctx.Builder(node->Pos())
4244+
.Callable("Coalesce")
4245+
.Callable(0, "Nth")
4246+
.Callable(0, "ListLast")
4247+
.Callable(0, "ListTakeWhile")
4248+
.Callable(0, "ListExtend")
4249+
.Callable(0, "AsList")
4250+
.List(0)
4251+
.Add(0, init)
4252+
.Add(1, init)
4253+
.Seal()
4254+
.Seal()
4255+
.Add(1, foldMap)
4256+
.Seal()
4257+
.Lambda(1)
4258+
.Param("x")
4259+
.Callable("HasItems")
4260+
.Callable(0, "Nth")
4261+
.Arg(0, "x")
4262+
.Atom(1, 1)
4263+
.Seal()
4264+
.Seal()
4265+
.Seal()
4266+
.Seal()
4267+
.Seal()
4268+
.Atom(1, 0)
4269+
.Seal()
4270+
.Callable(1, "AsList")
4271+
.Seal()
4272+
.Seal()
4273+
.Build();
4274+
}
4275+
41574276
} // namespace NYql

‎ydb/library/yql/core/common_opt/yql_co_pgselect.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,6 @@ TExprNode::TPtr ExpandPgGroupRef(const TExprNode::TPtr& node, TExprContext& ctx,
2525

2626
TExprNode::TPtr ExpandPgGrouping(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx);
2727

28+
TExprNode::TPtr ExpandPgIterate(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx);
29+
2830
} // namespace NYql

‎ydb/library/yql/core/common_opt/yql_co_simple1.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6257,6 +6257,8 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
62576257
};
62586258

62596259
map["PgSelect"] = &ExpandPgSelect;
6260+
map["PgIterate"] = &ExpandPgIterate;
6261+
map["PgIterateAll"] = &ExpandPgIterate;
62606262

62616263
map["PgLike"] = &ExpandPgLike;
62626264
map["PgILike"] = &ExpandPgLike;

‎ydb/library/yql/core/type_ann/type_ann_core.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12359,6 +12359,8 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1235912359
Functions["PgGrouping"] = &PgGroupingWrapper;
1236012360
Functions["PgGroupingSet"] = &PgGroupingSetWrapper;
1236112361
Functions["PgToRecord"] = &PgToRecordWrapper;
12362+
Functions["PgIterate"] = &PgIterateWrapper;
12363+
Functions["PgIterateAll"] = &PgIterateWrapper;
1236212364
Functions["StructUnion"] = &StructMergeWrapper;
1236312365
Functions["StructIntersection"] = &StructMergeWrapper;
1236412366
Functions["StructDifference"] = &StructMergeWrapper;

‎ydb/library/yql/core/type_ann/type_ann_pg.cpp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5815,5 +5815,40 @@ IGraphTransformer::TStatus PgToRecordWrapper(const TExprNode::TPtr& input, TExpr
58155815
return IGraphTransformer::TStatus::Ok;
58165816
}
58175817

5818+
IGraphTransformer::TStatus PgIterateWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
5819+
Y_UNUSED(output);
5820+
if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
5821+
return IGraphTransformer::TStatus::Error;
5822+
}
5823+
5824+
if (!EnsureListType(input->Head(), ctx.Expr)) {
5825+
return IGraphTransformer::TStatus::Error;
5826+
}
5827+
5828+
auto& lambda = input->ChildRef(1);
5829+
const auto status = ConvertToLambda(lambda, ctx.Expr, 1);
5830+
if (status.Level != IGraphTransformer::TStatus::Ok) {
5831+
return status;
5832+
}
5833+
5834+
if (!UpdateLambdaAllArgumentsTypes(lambda, { input->Head().GetTypeAnn() }, ctx.Expr)) {
5835+
return IGraphTransformer::TStatus::Error;
5836+
}
5837+
5838+
if (!lambda->GetTypeAnn()) {
5839+
return IGraphTransformer::TStatus::Repeat;
5840+
}
5841+
5842+
if (!IsSameAnnotation(*lambda->GetTypeAnn(), *input->Head().GetTypeAnn())) {
5843+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambda->Pos()), TStringBuilder() <<
5844+
"Mismatch of transform lambda return type and input type: " <<
5845+
*lambda->GetTypeAnn() << " != " << *input->Head().GetTypeAnn()));
5846+
return IGraphTransformer::TStatus::Error;
5847+
}
5848+
5849+
input->SetTypeAnn(input->Head().GetTypeAnn());
5850+
return IGraphTransformer::TStatus::Ok;
5851+
}
5852+
58185853
} // namespace NTypeAnnImpl
58195854
}

‎ydb/library/yql/core/type_ann/type_ann_pg.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ IGraphTransformer::TStatus PgGroupRefWrapper(const TExprNode::TPtr& input, TExpr
5757
IGraphTransformer::TStatus PgGroupingWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5858
IGraphTransformer::TStatus PgGroupingSetWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5959
IGraphTransformer::TStatus PgToRecordWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
60+
IGraphTransformer::TStatus PgIterateWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
6061

6162
} // namespace NTypeAnnImpl
6263
} // namespace NYql

‎ydb/library/yql/core/yql_type_annotation.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@ struct TTypeAnnotationContext: public TThrRefBase {
251251
ui32 EvaluateForLimit = 500;
252252
ui32 EvaluateParallelForLimit = 5000;
253253
ui32 EvaluateOrderByColumnLimit = 100;
254+
ui32 PgIterateLimit = 500;
254255
bool PullUpFlatMapOverJoin = true;
255256
bool FilterPushdownOverJoinOptionalSide = false;
256257
bool DeprecatedSQL = false;

‎ydb/library/yql/tests/sql/dq_file/part4/canondata/result.json‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2093,6 +2093,28 @@
20932093
}
20942094
],
20952095
"test.test[pg-pg_interval_literal-default.txt-Results]": [],
2096+
"test.test[pg-pg_iterate-default.txt-Analyze]": [
2097+
{
2098+
"checksum": "b2a2eb5d6b0a138ee924c128fc7738ef",
2099+
"size": 1331,
2100+
"uri": "https://{canondata_backend}/1916746/21c597d64388f6c41b4782746fbaf7f1b842bcda/resource.tar.gz#test.test_pg-pg_iterate-default.txt-Analyze_/plan.txt"
2101+
}
2102+
],
2103+
"test.test[pg-pg_iterate-default.txt-Debug]": [
2104+
{
2105+
"checksum": "99296cf8319e202900db3eb7f8dc163a",
2106+
"size": 2202,
2107+
"uri": "https://{canondata_backend}/1916746/21c597d64388f6c41b4782746fbaf7f1b842bcda/resource.tar.gz#test.test_pg-pg_iterate-default.txt-Debug_/opt.yql_patched"
2108+
}
2109+
],
2110+
"test.test[pg-pg_iterate-default.txt-Plan]": [
2111+
{
2112+
"checksum": "b2a2eb5d6b0a138ee924c128fc7738ef",
2113+
"size": 1331,
2114+
"uri": "https://{canondata_backend}/1916746/21c597d64388f6c41b4782746fbaf7f1b842bcda/resource.tar.gz#test.test_pg-pg_iterate-default.txt-Plan_/plan.txt"
2115+
}
2116+
],
2117+
"test.test[pg-pg_iterate-default.txt-Results]": [],
20962118
"test.test[pg-pg_like_opt-default.txt-Analyze]": [
20972119
{
20982120
"checksum": "b4dd508a329723c74293d80f0278c705",

‎ydb/library/yql/tests/sql/hybrid_file/part3/canondata/result.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1987,6 +1987,20 @@
19871987
"uri": "https://{canondata_backend}/1931696/fca86c589326e9bc05817a71a47f8b9d16219dcc/resource.tar.gz#test.test_pg-pg_corr_count-default.txt-Plan_/plan.txt"
19881988
}
19891989
],
1990+
"test.test[pg-pg_iterate-default.txt-Debug]": [
1991+
{
1992+
"checksum": "24c238402049dd496cfbcb0cf62a9d7d",
1993+
"size": 2201,
1994+
"uri": "https://{canondata_backend}/1942278/340f722a851e4412d2c35b434f3ee6113a9f7959/resource.tar.gz#test.test_pg-pg_iterate-default.txt-Debug_/opt.yql_patched"
1995+
}
1996+
],
1997+
"test.test[pg-pg_iterate-default.txt-Plan]": [
1998+
{
1999+
"checksum": "b2a2eb5d6b0a138ee924c128fc7738ef",
2000+
"size": 1331,
2001+
"uri": "https://{canondata_backend}/1942278/340f722a851e4412d2c35b434f3ee6113a9f7959/resource.tar.gz#test.test_pg-pg_iterate-default.txt-Plan_/plan.txt"
2002+
}
2003+
],
19902004
"test.test[pg-pg_types_window1-default.txt-Debug]": [
19912005
{
19922006
"checksum": "46e49d281a0b303ff2bb0e8409a16930",

‎ydb/library/yql/tests/sql/sql2yql/canondata/result.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11969,6 +11969,13 @@
1196911969
"uri": "https://{canondata_backend}/1925821/85af800a282697ba458d7b0567031a1eadd214a4/resource.tar.gz#test_sql2yql.test_pg-pg_interval_literal_/sql.yql"
1197011970
}
1197111971
],
11972+
"test_sql2yql.test[pg-pg_iterate]": [
11973+
{
11974+
"checksum": "776ef60a2ba5385372197ec39e44eb82",
11975+
"size": 4118,
11976+
"uri": "https://{canondata_backend}/1689644/58b94ee78c0b4d3f96a346070a311370cd242424/resource.tar.gz#test_sql2yql.test_pg-pg_iterate_/sql.yql"
11977+
}
11978+
],
1197211979
"test_sql2yql.test[pg-pg_like]": [
1197311980
{
1197411981
"checksum": "96d621bd2831dfda92b24d1eb2c96c0e",
@@ -29966,6 +29973,13 @@
2996629973
"uri": "https://{canondata_backend}/1880306/64654158d6bfb1289c66c626a8162239289559d0/resource.tar.gz#test_sql_format.test_pg-pg_in_dict_key_with_stable_pickle_/formatted.sql"
2996729974
}
2996829975
],
29976+
"test_sql_format.test[pg-pg_iterate]": [
29977+
{
29978+
"checksum": "1da98df588d09df6f6e91c46a5d06701",
29979+
"size": 715,
29980+
"uri": "https://{canondata_backend}/1689644/58b94ee78c0b4d3f96a346070a311370cd242424/resource.tar.gz#test_sql_format.test_pg-pg_iterate_/formatted.sql"
29981+
}
29982+
],
2996929983
"test_sql_format.test[pg-pg_type_from_oid]": [
2997029984
{
2997129985
"checksum": "ee7e7b037b636ac3ffedbb1f281c4995",

0 commit comments

Comments
 (0)