YQL-14728 in expression

ref:a756219b4eb959690a26aaca05a28434600ccaa6

7 files changed, 162 insertions, 0 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
index 277b9155fd..4147528ba1 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -1158,4 +1158,31 @@ TExprNode::TPtr ExpandPgLike(const TExprNode::TPtr& node, TExprContext& ctx, TOp
         .Build();
 }
 
+TExprNode::TPtr ExpandPgIn(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+    Y_UNUSED(optCtx);
+    return ctx.Builder(node->Pos())
+        .Callable("Fold")
+            .Add(0, node->ChildPtr(1))
+            .Callable(1, "PgConst")
+                .Atom(0, "false")
+                .Callable(1, "PgType")
+                    .Atom(0, "bool")
+                .Seal()
+            .Seal()
+            .Lambda(2)
+                .Param("item")
+                .Param("state")
+                .Callable("PgOr")
+                    .Arg(0, "state")
+                    .Callable(1, "PgOp")
+                        .Atom(0, "=")
+                        .Arg(1, "item")
+                        .Add(2, node->ChildPtr(0))
+                    .Seal()
+                .Seal()
+            .Seal()
+        .Seal()
+        .Build();
+}
+
 } // namespace NYql
diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.h b/ydb/library/yql/core/common_opt/yql_co_pgselect.h
index 17543841f3..24baabda5b 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.h
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.h
@@ -11,4 +11,6 @@ TExprNode::TPtr ExpandPositionalUnionAll(const TExprNode& node, const TVector<TC
 
 TExprNode::TPtr ExpandPgLike(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx);
 
+TExprNode::TPtr ExpandPgIn(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx);
+
 } // namespace NYql
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index 5bf411f8b7..97873f16c6 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -6365,6 +6365,8 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
     map["PgLike"] = &ExpandPgLike;
     map["PgILike"] = &ExpandPgLike;
 
+    map["PgIn"] = &ExpandPgIn;
+
     map["SqlColumnOrType"] = map["SqlPlainColumnOrType"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& /*optCtx*/) {
         YQL_CLOG(DEBUG, Core) << "Decay of never inspected " << node->Content();
         return ctx.NewCallable(node->Pos(), "Error", { ExpandType(node->Pos(), *node->GetTypeAnn(), ctx) });
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 76af8fd3df..3a50d6b1bc 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -11286,6 +11286,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["PgTypeMod"] = &PgTypeModWrapper;
         Functions["PgLike"] = &PgLikeWrapper;
         Functions["PgILike"] = &PgLikeWrapper;
+        Functions["PgIn"] = &PgInWrapper;
 
         Functions["AutoDemuxList"] = &AutoDemuxListWrapper;
         Functions["AggrCountInit"] = &AggrCountInitWrapper;
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
index 7a5e20767a..fb95bc8acd 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -2741,5 +2741,79 @@ IGraphTransformer::TStatus PgLikeWrapper(const TExprNode::TPtr& input, TExprNode
     return IGraphTransformer::TStatus::Ok;
 }
 
+IGraphTransformer::TStatus PgInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+    if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    auto inputType = input->Child(0)->GetTypeAnn();
+    ui32 inputTypePg;
+    bool convertToPg;
+    if (!ExtractPgType(inputType, inputTypePg, convertToPg, input->Child(0)->Pos(), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (convertToPg) {
+        input->ChildRef(0) = ctx.Expr.NewCallable(input->Child(0)->Pos(), "ToPg", { input->ChildPtr(0) });
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    auto listType = input->Child(1)->GetTypeAnn();
+    if (listType && listType->GetKind() == ETypeAnnotationKind::EmptyList) {
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), "IN expects at least one element"));
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (!EnsureListType(*input->Child(1), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    auto listItemType = listType->Cast<TListExprType>()->GetItemType();
+    ui32 itemTypePg;
+    if (!ExtractPgType(listItemType, itemTypePg, convertToPg, input->Child(1)->Pos(), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (convertToPg) {
+        output = ctx.Expr.Builder(input->Pos())
+            .Callable("PgIn")
+                .Add(0, input->ChildPtr(0))
+                .Callable(1, "Map")
+                    .Add(0, input->ChildPtr(1))
+                    .Lambda(1)
+                        .Param("x")
+                        .Callable("ToPg")
+                            .Arg(0, "x")
+                        .Seal()
+                    .Seal()
+                .Seal()
+            .Seal()
+            .Build();
+
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    if (itemTypePg && inputTypePg && itemTypePg != inputTypePg) {
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Mismatch of types in IN expressions: " <<
+            NPg::LookupType(inputTypePg).Name << " is not equal to " << NPg::LookupType(itemTypePg).Name));
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (!listItemType->IsEquatable()) {
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Cannot compare items of type: " << NPg::LookupType(itemTypePg).Name));
+    }
+
+    if (!inputType->IsEquatable()) {
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Cannot compare items of type: " << NPg::LookupType(inputTypePg).Name));
+    }
+
+    auto result = ctx.Expr.MakeType<TPgExprType>(NPg::LookupType("bool").TypeId);
+    input->SetTypeAnn(result);
+    return IGraphTransformer::TStatus::Ok;
+}
+
 } // namespace NTypeAnnImpl
 }
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.h b/ydb/library/yql/core/type_ann/type_ann_pg.h
index 7f19e3a7bb..08fdd3296e 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.h
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.h
@@ -33,6 +33,7 @@ IGraphTransformer::TStatus PgSelectWrapper(const TExprNode::TPtr& input, TExprNo
 IGraphTransformer::TStatus PgArrayWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgTypeModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgLikeWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+IGraphTransformer::TStatus PgInWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 
 } // namespace NTypeAnnImpl
 } // namespace NYql
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index 5fcfaa907a..cba5991b87 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -1827,6 +1827,59 @@ public:
         return ret;
     }
 
+    TAstNode* ParseAExprIn(const A_Expr* value, const TExprSettings& settings) {
+        if (ListLength(value->name) != 1) {
+            AddError(TStringBuilder() << "Unsupported count of names: " << ListLength(value->name));
+            return nullptr;
+        }
+
+        auto nameNode = ListNodeNth(value->name, 0);
+        if (NodeTag(nameNode) != T_String) {
+            NodeNotImplemented(value, nameNode);
+            return nullptr;
+        }
+
+        auto op = StrVal(nameNode);
+        if (op != "=" && op != "<>") {
+            AddError(TStringBuilder() << "Unsupported operation: " << op);
+            return nullptr;
+        }
+
+        if (!value->lexpr || !value->rexpr) {
+            AddError("Missing operands");
+            return nullptr;
+        }
+
+        auto lhs = ParseExpr(value->lexpr, settings);
+        if (!lhs) {
+            return nullptr;
+        }
+
+        if (NodeTag(value->rexpr) != T_List) {
+            NodeNotImplemented(value, value->rexpr);
+            return nullptr;
+        }
+
+        auto lst = CAST_NODE(List, value->rexpr);
+        TVector<TAstNode*> listItems;
+        listItems.push_back(A("AsList"));
+        for (int item = 0; item < ListLength(lst); ++item) {
+            auto cell = ParseExpr(ListNodeNth(lst, item), settings);
+            if (!cell) {
+                return nullptr;
+            }
+
+            listItems.push_back(cell);
+        }
+
+        auto ret = L(A("PgIn"), lhs, VL(listItems.data(), listItems.size()));
+        if (op[0] == '<') {
+            ret = L(A("PgNot"), ret);
+        }
+
+        return ret;
+    }
+
     TAstNode* ParseAExprBetween(const A_Expr* value, const TExprSettings& settings) {
         if (!value->lexpr || !value->rexpr) {
             AddError("Missing operands");
@@ -1898,6 +1951,8 @@ public:
         case AEXPR_LIKE:
         case AEXPR_ILIKE:
             return ParseAExprLike(value, settings, value->kind == AEXPR_ILIKE);
+        case AEXPR_IN:
+            return ParseAExprIn(value, settings);
         case AEXPR_BETWEEN:
         case AEXPR_NOT_BETWEEN:
         case AEXPR_BETWEEN_SYM: