Add some missed optimizers and a little refactoring.

4 files changed, 49 insertions, 36 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
index 516bdda32cc..9b1984e885d 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
@@ -1182,12 +1182,9 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
             TSet<TStringBuf> usedFields;
             if (HaveFieldsSubset(body, arg, usedFields, *optCtx.ParentsMap)) {
                 YQL_CLOG(DEBUG, Core) << "FieldsSubset in " << node->Content() << " over " << self.Input().Ref().Content();
-                TSet<TString> fields;
-                for (auto& x : usedFields) {
-                    fields.emplace(TString(x));
-                }
 
                 TExprNode::TListType filteredInputs;
+                filteredInputs.reserve(self.Input().Ref().ChildrenSize());
                 for (ui32 index = 0; index < self.Input().Ref().ChildrenSize(); ++index) {
                     auto x = self.Input().Ref().ChildPtr(index);
                     if (!self.Input().Maybe<TCoExtendBase>() && index > 0) {
@@ -1195,7 +1192,7 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
                         continue;
                     }
 
-                    filteredInputs.push_back(FilterByFields(node->Pos(), x, fields, ctx, false));
+                    filteredInputs.push_back(FilterByFields(node->Pos(), x, usedFields, ctx, false));
                 }
 
                 auto newInput = ctx.ChangeChildren(self.Input().Ref(), std::move(filteredInputs));
@@ -1284,12 +1281,8 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
 
         TExprNode::TPtr newInput;
         if (self.Input().Ref().IsCallable("Take") || self.Input().Ref().IsCallable("Skip") || self.Input().Maybe<TCoExtendBase>()) {
-            TSet<TString> fields;
-            for (auto& x : usedFields) {
-                fields.emplace(TString(x));
-            }
-
             TExprNode::TListType filteredInputs;
+            filteredInputs.reserve(self.Input().Ref().ChildrenSize());
             for (ui32 index = 0; index < self.Input().Ref().ChildrenSize(); ++index) {
                 auto x = self.Input().Ref().ChildPtr(index);
                 if (!self.Input().Maybe<TCoExtendBase>() && index > 0) {
@@ -1297,7 +1290,7 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
                     continue;
                 }
 
-                filteredInputs.push_back(FilterByFields(node->Pos(), x, fields, ctx, false));
+                filteredInputs.push_back(FilterByFields(node->Pos(), x, usedFields, ctx, false));
             }
 
             YQL_CLOG(DEBUG, Core) << "FieldsSubset in " << node->Content() << " over " << self.Input().Ref().Content();
@@ -1796,7 +1789,7 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
 
         auto ret = AggregateSubsetFieldsAnalyzer(self, ctx, *optCtx.ParentsMap);
         if (ret != node) {
-            YQL_CLOG(DEBUG, Core) << "AggregateSubsetFieldsAnalyzer";
+            YQL_CLOG(DEBUG, Core) << node->Content() << "SubsetFieldsAnalyzer";
             return ret;
         }
 
@@ -1806,21 +1799,21 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
     map["CalcOverWindow"] = map["CalcOverSessionWindow"] = map["CalcOverWindowGroup"] =
         [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx)
     {
-        if (!optCtx.IsSingleUsage(*node->Child(0))) {
+        if (!optCtx.IsSingleUsage(node->Head())) {
             return node;
         }
 
-        if (!node->Child(0)->IsCallable({"CalcOverWindow", "CalcOverSessionWindow", "CalcOverWindowGroup"})) {
+        if (!node->Head().IsCallable({"CalcOverWindow", "CalcOverSessionWindow", "CalcOverWindowGroup"})) {
             return node;
         }
 
         TExprNodeList parentCalcs = ExtractCalcsOverWindow(node, ctx);
-        TExprNodeList calcs = ExtractCalcsOverWindow(node->ChildPtr(0), ctx);
+        TExprNodeList calcs = ExtractCalcsOverWindow(node->HeadPtr(), ctx);
         calcs.insert(calcs.end(), parentCalcs.begin(), parentCalcs.end());
 
-        YQL_CLOG(DEBUG, Core) << "Fuse nested CalcOverWindow/CalcOverSessionWindow/CalcOverWindowGroup";
+        YQL_CLOG(DEBUG, Core) << "Fuse nested " << node->Content() << " and " << node->Head().Content();
 
-        return RebuildCalcOverWindowGroup(node->Child(0)->Pos(), node->Child(0)->ChildPtr(0), calcs, ctx);
+        return RebuildCalcOverWindowGroup(node->Head().Pos(), node->Head().HeadPtr(), calcs, ctx);
     };
 
     map[TCoCondense::CallableName()] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
@@ -1955,6 +1948,40 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
         }
         return node;
     };
+
+    map[TCoSqueezeToDict::CallableName()] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        const TCoSqueezeToDict self(node);
+        if (!optCtx.IsSingleUsage(self.Stream().Ref())) {
+            return node;
+        }
+
+        std::map<std::string_view, TExprNode::TPtr> usedFields;
+        if (HaveFieldsSubset(self.KeySelector().Body().Ptr(), self.KeySelector().Args().Arg(0).Ref(), usedFields, *optCtx.ParentsMap, false)
+            && !usedFields.empty()
+            && HaveFieldsSubset(self.PayloadSelector().Body().Ptr(), self.PayloadSelector().Args().Arg(0).Ref(), usedFields, *optCtx.ParentsMap, false)
+            && !usedFields.empty()
+            && usedFields.size() < GetSeqItemType(self.Stream().Ref().GetTypeAnn())->Cast<TStructExprType>()->GetSize())
+        {
+            TExprNode::TListType fields;
+            fields.reserve(usedFields.size());
+            std::transform(usedFields.begin(), usedFields.end(), std::back_inserter(fields),
+                [](std::pair<const std::string_view, TExprNode::TPtr>& item){ return std::move(item.second); });
+
+            YQL_CLOG(DEBUG, Core) << node->Content() << "SubsetFields";
+            return Build<TCoSqueezeToDict>(ctx, node->Pos())
+                .Stream<TCoExtractMembers>()
+                    .Input(self.Stream())
+                    .Members()
+                        .Add(std::move(fields))
+                    .Build()
+                .Build()
+                .KeySelector(ctx.DeepCopyLambda(self.KeySelector().Ref()))
+                .PayloadSelector(ctx.DeepCopyLambda(self.PayloadSelector().Ref()))
+                .Settings(self.Settings())
+                .Done().Ptr();
+        }
+        return node;
+    };
 }
 
 }
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 a0f5d4e15bd..175bac79505 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -3629,14 +3629,9 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
 
         if (node->Head().IsCallable({"Just", "AsList"})) {
             YQL_CLOG(DEBUG, Core) << "Move " << node->Content() << " over " << node->Head().Content();
-            TSet<TString> fields;
-            node->Tail().ForEachChild([&fields](const TExprNode& child) {
-                fields.emplace(child.Content());
-            });
-
             auto args = node->Head().ChildrenList();
             for (auto& arg : args) {
-                arg = FilterByFields(node->Pos(), arg, fields, ctx, true);
+                arg = ctx.NewCallable(node->Pos(), "FilterMembers", {std::move(arg), node->TailPtr()});
             }
 
             return ctx.ChangeChildren(node->Head(), std::move(args));
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 20770b952b8..f2ad31cea4c 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -115,6 +115,7 @@ public:
         Functions["RemoveSystemMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
         Functions["RemovePrefixMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
         Functions["SelectMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
+        Functions["FilterMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
         Functions["CastStruct"] = &TCallableConstraintTransformer::SelectMembersWrap;
         Functions["SafeCast"] = &TCallableConstraintTransformer::SelectMembersWrap<true>;
         Functions["StrictCast"] = &TCallableConstraintTransformer::SelectMembersWrap<true>;
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index 67c5fbf8b6f..fe94ead345e 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -307,23 +307,13 @@ TExprNode::TPtr AddMembersUsedInside(const TExprNode::TPtr& start, const TExprNo
 template<class TFieldsSet>
 TExprNode::TPtr FilterByFields(TPositionHandle position, const TExprNode::TPtr& input, const TFieldsSet& subsetFields,
     TExprContext& ctx, bool singleValue) {
-    if (singleValue) {
-        TExprNode::TListType structItems;
-        for (auto field : subsetFields) {
-            auto name = ctx.NewAtom(position, field);
-            auto member = ctx.NewCallable(position, "Member", { input , name });
-            structItems.push_back(ctx.NewList(position, { name, member }));
-        }
-
-        return ctx.NewCallable(position, "AsStruct", std::move(structItems));
-    }
-
     TExprNode::TListType fields;
-    for (auto& x : subsetFields) {
+    fields.reserve(subsetFields.size());
+    for (const auto& x : subsetFields) {
         fields.emplace_back(ctx.NewAtom(position, x));
     }
 
-    return ctx.NewCallable(position, "ExtractMembers", { input, ctx.NewList(position, std::move(fields)) });
+    return ctx.NewCallable(position, singleValue ? "FilterMembers" : "ExtractMembers", { input, ctx.NewList(position, std::move(fields)) });
 }
 
 template TExprNode::TPtr FilterByFields(TPositionHandle position, const TExprNode::TPtr& input, const TSet<TStringBuf>& subsetFields, TExprContext& ctx, bool singleValue);