Skip to content

Commit 942f92c

Browse files
marsaly79marsaly79
committed
Implemented IN with varying LHS and RHS types
1 parent e264ef6 commit 942f92c

File tree

6 files changed

+168
-5
lines changed

6 files changed

+168
-5
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12070,6 +12070,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1207012070
Functions["PgLike"] = &PgLikeWrapper;
1207112071
Functions["PgILike"] = &PgLikeWrapper;
1207212072
Functions["PgIn"] = &PgInWrapper;
12073+
Functions["PgMixedIn"] = &PgMixedInWrapper;
1207312074
Functions["PgBetween"] = &PgBetweenWrapper;
1207412075
Functions["PgBetweenSym"] = &PgBetweenWrapper;
1207512076
Functions["PgSubLink"] = &PgSubLinkWrapper;

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

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5285,6 +5285,142 @@ IGraphTransformer::TStatus PgLikeWrapper(const TExprNode::TPtr& input, TExprNode
52855285
return IGraphTransformer::TStatus::Ok;
52865286
}
52875287

5288+
TExprNodePtr BuildPgIn(TExprNodeList&& args, ui32 targetTypeId, bool castRequired, TContext& ctx) {
5289+
auto lhs = args[0];
5290+
if (castRequired) {
5291+
const auto lhsType = lhs->GetTypeAnn()->Cast<TPgExprType>();
5292+
if (lhsType->GetId() != targetTypeId) {
5293+
lhs = WrapWithPgCast(std::move(lhs), targetTypeId, ctx.Expr);
5294+
}
5295+
}
5296+
std::swap(args[0], args.back());
5297+
args.pop_back();
5298+
5299+
if (castRequired) {
5300+
for (size_t i = 0; i < args.size(); ++i) {
5301+
const auto argType = args[i]->GetTypeAnn()->Cast<TPgExprType>();
5302+
if (argType->GetId() != targetTypeId) {
5303+
args[i] = WrapWithPgCast(std::move(args[i]), targetTypeId, ctx.Expr);
5304+
}
5305+
}
5306+
}
5307+
5308+
const auto rhs = ctx.Expr.Builder(args[1]->Pos())
5309+
.Callable("AsList")
5310+
.Add(std::move(args))
5311+
.Seal()
5312+
.Build();
5313+
5314+
return ctx.Expr.Builder(lhs->Pos())
5315+
.Callable("PgIn")
5316+
.Add(0, lhs)
5317+
.Add(1, rhs)
5318+
.Seal()
5319+
.Build();
5320+
}
5321+
5322+
IGraphTransformer::TStatus PgMixedInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
5323+
if (!EnsureMinArgsCount(*input, 2, ctx.Expr)) {
5324+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), "IN expects at least one element"));
5325+
return IGraphTransformer::TStatus::Error;
5326+
}
5327+
5328+
TVector<ui32> pgTypes(input->ChildrenSize());
5329+
{
5330+
TExprNodeList convertedChildren;
5331+
convertedChildren.reserve(input->ChildrenSize());
5332+
bool convertionRequired = false;
5333+
bool hasConvertions = false;
5334+
for (size_t i = 0; i < input->ChildrenSize(); ++i) {
5335+
const auto child = input->Child(i);
5336+
if (!ExtractPgType(child->GetTypeAnn(), pgTypes[i], convertionRequired, child->Pos(), ctx.Expr)) {
5337+
return IGraphTransformer::TStatus::Error;
5338+
}
5339+
if (convertionRequired) {
5340+
hasConvertions = true;
5341+
5342+
auto convertedChild = ctx.Expr.Builder(child->Pos())
5343+
.Callable("ToPg")
5344+
.Add(0, child)
5345+
.Seal()
5346+
.Build();
5347+
convertedChildren.push_back(std::move(convertedChild));
5348+
} else {
5349+
convertedChildren.push_back(std::move(child));
5350+
}
5351+
}
5352+
if (hasConvertions) {
5353+
output = ctx.Expr.Builder(input->Pos())
5354+
.Callable("PgMixedIn")
5355+
.Add(std::move(convertedChildren))
5356+
.Seal()
5357+
.Build();
5358+
5359+
return IGraphTransformer::TStatus::Repeat;
5360+
}
5361+
}
5362+
5363+
auto posGetter = [&input, &ctx](size_t i) {
5364+
return ctx.Expr.GetPosition(input->Child(i)->Pos());
5365+
};
5366+
5367+
struct TPgListCommonTypeConversion {
5368+
bool conversionRequired = false;
5369+
ui32 targetType = 0;
5370+
TExprNodeList items;
5371+
};
5372+
5373+
bool castRequired = false;
5374+
const NPg::TTypeDesc* commonType;
5375+
if (NPg::LookupCommonType(pgTypes, posGetter, commonType, castRequired))
5376+
{
5377+
THashMap<ui32, TPgListCommonTypeConversion> elemsByType;
5378+
for (size_t i = 1; i < input->ChildrenSize(); ++i) {
5379+
auto& elemsOfType = elemsByType[pgTypes[i]];
5380+
if (elemsOfType.items.empty()) {
5381+
const NPg::TTypeDesc* elemCommonType;
5382+
if (const auto issue = NPg::LookupCommonType({pgTypes[0], pgTypes[i]},
5383+
posGetter, elemCommonType, elemsOfType.conversionRequired))
5384+
{
5385+
ctx.Expr.AddError(*issue);
5386+
return IGraphTransformer::TStatus::Error;
5387+
}
5388+
elemsOfType.targetType = elemCommonType->TypeId;
5389+
// Clone IN's lhs expression as the leftmost elem in the new type-associated list
5390+
elemsOfType.items.push_back(ctx.Expr.ShallowCopy(input->Head()));
5391+
elemsOfType.items.back()->SetTypeAnn(input->Head().GetTypeAnn());
5392+
}
5393+
elemsOfType.items.push_back(input->Child(i));
5394+
}
5395+
TExprNodeList orClausesOfIn;
5396+
orClausesOfIn.reserve(elemsByType.size());
5397+
5398+
for (auto& elemsOfType: elemsByType) {
5399+
auto& conversion = elemsOfType.second;
5400+
orClausesOfIn.push_back(BuildPgIn(std::move(conversion.items), conversion.targetType, conversion.conversionRequired, ctx));
5401+
}
5402+
output = ctx.Expr.Builder(input->Pos())
5403+
.Callable("Fold")
5404+
.Callable(0, "AsList")
5405+
.Add(std::move(orClausesOfIn))
5406+
.Seal()
5407+
.Add(1, MakePgBool(input->Pos(), false, ctx.Expr))
5408+
.Lambda(2)
5409+
.Param("item")
5410+
.Param("state")
5411+
.Callable("PgOr")
5412+
.Arg(0, "state")
5413+
.Arg(1, "item")
5414+
.Seal()
5415+
.Seal()
5416+
.Seal()
5417+
.Build();
5418+
} else {
5419+
output = BuildPgIn(std::move(input->ChildrenList()), commonType->TypeId, castRequired, ctx);
5420+
}
5421+
return IGraphTransformer::TStatus::Repeat;
5422+
}
5423+
52885424
IGraphTransformer::TStatus PgInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
52895425
if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
52905426
return IGraphTransformer::TStatus::Error;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ IGraphTransformer::TStatus PgArrayWrapper(const TExprNode::TPtr& input, TExprNod
5050
IGraphTransformer::TStatus PgTypeModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5151
IGraphTransformer::TStatus PgLikeWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5252
IGraphTransformer::TStatus PgInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
53+
IGraphTransformer::TStatus PgMixedInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5354
IGraphTransformer::TStatus PgBetweenWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5455
IGraphTransformer::TStatus PgSubLinkWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
5556
IGraphTransformer::TStatus PgGroupRefWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

‎ydb/library/yql/sql/pg/pg_sql.cpp‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4603,18 +4603,21 @@ class TConverter : public IPGParseEvents {
46034603
}
46044604

46054605
auto lst = CAST_NODE(List, value->rexpr);
4606-
TVector<TAstNode*> listItems;
4607-
listItems.push_back(A("AsList"));
4606+
4607+
TVector<TAstNode*> children;
4608+
children.reserve(2 + ListLength(lst));
4609+
4610+
children.push_back(A("PgMixedIn"));
4611+
children.push_back(lhs);
46084612
for (int item = 0; item < ListLength(lst); ++item) {
46094613
auto cell = ParseExpr(ListNodeNth(lst, item), settings);
46104614
if (!cell) {
46114615
return nullptr;
46124616
}
4613-
4614-
listItems.push_back(cell);
4617+
children.push_back(cell);
46154618
}
46164619

4617-
auto ret = L(A("PgIn"), lhs, VL(listItems.data(), listItems.size()));
4620+
auto ret = VL(children.data(), children.size());
46184621
if (op[0] == '<') {
46194622
ret = L(A("PgNot"), ret);
46204623
}

‎ydb/library/yql/tests/sql/suites/pg/in_mixed.cfg‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
in Input input_name.txt

‎ydb/library/yql/tests/sql/suites/pg/in_mixed.sql‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
--!syntax_pg
2+
3+
-- simple case
4+
select 1 in ('1', '2');
5+
6+
-- splits into IN over numerics and IN over text
7+
select '1' in (0, 0.0, 1, '1'::text);
8+
9+
-- mixture of types in rhs, expression in lhs
10+
select (c::int + 1) in (1, 2.9, '4') from (values ('0'), ('1')) as t(c);
11+
12+
-- arrays support
13+
select array[1, 2] in (array[2, 4], array[1, 2]);
14+
15+
-- NULL in rhs
16+
select 1 in (0, NULL);
17+
select 1 in (1, NULL);
18+
19+
-- non-PG types handling
20+
select index in ('2', '4') from "Input";
21+

0 commit comments

Comments
 (0)