ListUniqStable

YQL-9448

7 files changed, 197 insertions, 27 deletions

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 68a1155ac4a..be6d4eb9e02 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -3706,6 +3706,156 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
     map["Lookup"] = std::bind(&OptimizeContains<false, true>, _1, _2);
     map["Contains"] = std::bind(&OptimizeContains<false>, _1, _2);
     map["ListHas"] = std::bind(&OptimizeContains<true>, _1, _2);
+    map["ListUniq"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext&) {
+        return ctx.Builder(node->Pos())
+            .Callable("DictKeys")
+                .Callable(0, "ToDict")
+                    .Add(0, node->HeadPtr())
+                    .Lambda(1)
+                        .Param("item")
+                        .Arg("item")
+                    .Seal()
+                    .Lambda(2)
+                        .Param("item")
+                        .Callable("Void")
+                        .Seal()
+                    .Seal()
+                    .List(3)
+                        .Atom(0, "Hashed")
+                        .Atom(1, "One")
+                    .Seal()
+                .Seal()
+            .Seal()
+            .Build();
+    };
+    map["ListUniqStable"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext&) {
+        const TTypeAnnotationNode* itemType = node->Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType();
+        auto expandedItemType = ExpandType(node->Pos(), *itemType, ctx);
+        auto setCreate = ctx.Builder(node->Pos())
+            .Callable("Udf")
+                .Atom(0, "Set.Create")
+                .Callable(1, "Void").Seal()
+                .Callable(2, "TupleType")
+                    .Callable(0, "TupleType")
+                        .Add(0, expandedItemType)
+                        .Callable(1, "DataType")
+                            .Atom(0, "Uint32", TNodeFlags::Default)
+                        .Seal()
+                    .Seal()
+                    .Callable(1, "StructType").Seal()
+                    .Add(2, expandedItemType)
+                .Seal()
+            .Seal()
+            .Build();
+
+        auto resourceType = ctx.Builder(node->Pos())
+            .Callable("TypeOf")
+                .Callable(0, "Apply")
+                    .Add(0, setCreate)
+                    .Callable(1, "InstanceOf")
+                        .Add(0, expandedItemType)
+                    .Seal()
+                    .Callable(2, "Uint32")
+                        .Atom(0, 0u)
+                    .Seal()
+                .Seal()
+            .Seal()
+            .Build();
+        auto setAddValue = ctx.Builder(node->Pos())
+            .Callable("Udf")
+                .Atom(0, "Set.AddValue")
+                .Callable(1, "Void").Seal()
+                .Callable(2, "TupleType")
+                    .Callable(0, "TupleType")
+                        .Add(0, resourceType)
+                        .Add(1, expandedItemType)
+                    .Seal()
+                    .Callable(1, "StructType").Seal()
+                    .Add(2, expandedItemType)
+                .Seal()
+            .Seal()
+            .Build();
+
+        auto setWasChanged = ctx.Builder(node->Pos())
+            .Callable("Udf")
+                .Atom(0, "Set.WasChanged")
+                .Callable(1, "Void").Seal()
+                .Callable(2, "TupleType")
+                    .Callable(0, "TupleType")
+                        .Add(0, resourceType)
+                    .Seal()
+                    .Callable(1, "StructType").Seal()
+                    .Add(2, expandedItemType)
+                .Seal()
+            .Seal()
+            .Build();
+
+        auto handlerLambda = ctx.Builder(node->Pos())
+            .Lambda()
+                .Param("updatedSet")
+                .Param("value")
+                .List(0)
+                    .Callable(0, "If")
+                        .Callable(0, "Apply")
+                            .Add(0, setWasChanged)
+                            .Arg(1, "updatedSet")
+                        .Seal()
+                        .Callable(1, "Just")
+                            .Arg(0, "value")
+                        .Seal()
+                        .Callable(2, "Nothing")
+                            .Callable(0, "OptionalType")
+                                .Add(0, expandedItemType)
+                            .Seal()
+                        .Seal()
+                    .Seal()
+                    .Arg(1, "updatedSet")
+                .Seal()
+            .Seal()
+            .Build();
+
+        return ctx.Builder(node->Pos())
+            .Callable("FlatMap")
+                .Callable(0, "Fold1Map")
+                    .Add(0, node->HeadPtr())
+                    .Lambda(1)
+                        .Param("item")
+                        .List(0)
+                            .Callable(0, "Just")
+                                .Arg(0, "item")
+                            .Seal()
+                            .Callable(1, "Apply")
+                                .Add(0, setCreate)
+                                .Arg(1, "item")
+                                .Callable(2, "Uint32")
+                                    .Atom(0, 0u)
+                                .Seal()
+                            .Seal()
+                        .Seal()
+                    .Seal()
+                    .Lambda(2)
+                        .Param("item")
+                        .Param("cur_set")
+                        .Apply(0, handlerLambda)
+                            .With(0)
+                                .Callable("Apply")
+                                    .Add(0, setAddValue)
+                                    .Arg(1, "cur_set")
+                                    .Arg(2, "item")
+                                .Seal()
+                            .Done()
+                            .With(1, "item")
+                        .Seal()
+                    .Seal()
+                .Seal()
+                .Lambda(1)
+                    .Param("item")
+                    .Arg(0, "item")
+                .Seal()
+            .Seal()
+            .Build();
+    };
+
     map["SqlIn"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext&) {
         auto collection = node->HeadPtr();
         auto lookup = node->ChildPtr(1);
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 7ccb328716c..00c767a7813 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -12020,6 +12020,8 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["ListAny"] = &ListAllAnyWrapper<false>;
         Functions["ListNotNull"] = &ListNotNullWrapper;
         Functions["ListFlatten"] = &ListFlattenWrapper;
+        Functions["ListUniq"] = &ListUniqWrapper;
+        Functions["ListUniqStable"] = &ListUniqWrapper;
 
         Functions["ExpandMap"] = &ExpandMapWrapper;
         Functions["WideMap"] = &WideMapWrapper;
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
index 6e6eee3a13f..f6764959711 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -7095,6 +7095,48 @@ namespace {
         return IGraphTransformer::TStatus::Repeat;
     }
 
+    IGraphTransformer::TStatus ListUniqWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureComputable(input->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (IsNull(input->Head())) {
+            output = input->HeadPtr();
+            return IGraphTransformer::TStatus::Repeat;
+        }
+
+        auto type = input->Head().GetTypeAnn();
+        if (type->GetKind() == ETypeAnnotationKind::Optional) {
+            type = type->Cast<TOptionalExprType>()->GetItemType();
+        }
+
+        if (type->GetKind() != ETypeAnnotationKind::List && type->GetKind() != ETypeAnnotationKind::EmptyList) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Head().Pos()), TStringBuilder()
+                << "Expected (empty) list or optional of (empty) list, but got: " << *input->Head().GetTypeAnn()));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (type->GetKind() == ETypeAnnotationKind::EmptyList) {
+            output = input->HeadPtr();
+            return IGraphTransformer::TStatus::Repeat;
+        }
+
+        auto itemType = type->Cast<TListExprType>()->GetItemType();
+
+        if (!itemType->IsHashable() || !itemType->IsEquatable()) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() << "Expected hashable and equatable type, but got: " << *itemType));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        input->SetTypeAnn(input->Head().GetTypeAnn());
+        
+        return IGraphTransformer::TStatus::Ok;
+    }
+
     IGraphTransformer::TStatus ListFlattenWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.h b/ydb/library/yql/core/type_ann/type_ann_list.h
index 5512c4dcbe7..50ac73b554e 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.h
+++ b/ydb/library/yql/core/type_ann/type_ann_list.h
@@ -128,6 +128,7 @@ namespace NTypeAnnImpl {
     template <bool AllOrAny>
     IGraphTransformer::TStatus ListAllAnyWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListNotNullWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+    IGraphTransformer::TStatus ListUniqWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListFlattenWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus IterableWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus SqueezeToListWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index af34005bd3d..dc67b0a4e44 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -2867,7 +2867,8 @@ struct TBuiltinFuncData {
             {"listavg", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListAvg", 1, 1)},
             {"listconcat", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListConcat", 1, 2)},
             {"listextract", BuildSimpleBuiltinFactoryCallback<TListExtractBuiltin>()},
-            {"listuniq", BuildSimpleBuiltinFactoryCallback<TListUniqBuiltin>()},
+            {"listuniq", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListUniq", 1, 1)},
+            {"listuniqstable", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListUniqStable", 1, 1)},
             {"listcreate", BuildSimpleBuiltinFactoryCallback<TListCreateBuiltin>()},
             {"listfromrange", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListFromRange", 2, 3) },
             {"listreplicate", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("Replicate", 2, 2) },
diff --git a/ydb/library/yql/sql/v1/list_builtin.cpp b/ydb/library/yql/sql/v1/list_builtin.cpp
index bc4edc2e4c1..5a264db97c2 100644
--- a/ydb/library/yql/sql/v1/list_builtin.cpp
+++ b/ydb/library/yql/sql/v1/list_builtin.cpp
@@ -85,18 +85,6 @@ TNodePtr TListFilterBuiltin::GetFilterLambda() {
     return BuildLambda(Pos, Y("item"), Y("Coalesce", Y("Apply", Args[1], "item"), Y("Bool", Q("false"))));
 }
 
-bool TListUniqBuiltin::DoInit(TContext& ctx, ISource* src) {
-    if (Args.size() != 1) {
-        ctx.Error(Pos) << OpName << " requires only one parameter";
-        return false;
-    }
-    if (!Args[0]->Init(ctx, src)) {
-        return false;
-    }
-    Node = Y("DictKeys", Y("ToDict", Args[0], GetIdentityLambda(), BuildLambda(Pos, Y("item"), Y("Void")), Q(Y(Q("Hashed"), Q("One")))));
-    return true;
-}
-
 bool TListCreateBuiltin::DoInit(TContext& ctx, ISource* src) {
     if (Args.size() != 1) {
         ctx.Error(Pos) << OpName << " requires only one parameter";
diff --git a/ydb/library/yql/sql/v1/list_builtin.h b/ydb/library/yql/sql/v1/list_builtin.h
index a8077b464a2..2a6e0112b15 100644
--- a/ydb/library/yql/sql/v1/list_builtin.h
+++ b/ydb/library/yql/sql/v1/list_builtin.h
@@ -112,20 +112,6 @@ protected:
     virtual TNodePtr GetFilterLambda();
 };
 
-class TListUniqBuiltin final: public TListBuiltin {
-public:
-    TListUniqBuiltin(TPosition pos,
-                     const TVector<TNodePtr>& args)
-        : TListBuiltin(pos, "ListUniq", args)
-    {}
-
-    bool DoInit(TContext& ctx, ISource* src) override;
-
-    TNodePtr DoClone() const final {
-        return new TListUniqBuiltin(Pos, CloneContainer(Args));
-    }
-};
-
 class TListCreateBuiltin final: public TListBuiltin {
 public:
     TListCreateBuiltin(TPosition pos,