[YQL-14389] Add MapNext optimizers

ref:2c84d356de4b420dea6fd5cf2b62c2c6ab307e70

7 files changed, 101 insertions, 2 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
index 215a1e9f69..35b24e9439 100644
--- a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
@@ -4,6 +4,7 @@
 #include <ydb/library/yql/core/yql_join.h>
 #include <ydb/library/yql/core/yql_opt_utils.h>
 #include <ydb/library/yql/core/yql_opt_window.h>
+#include <ydb/library/yql/core/yql_type_helpers.h>
 
 #include <ydb/library/yql/utils/log/log.h>
 
@@ -794,4 +795,30 @@ TExprNode::TPtr ApplyExtractMembersToCollect(const TExprNode::TPtr& node, const
         .Done().Ptr();
 }
 
+TExprNode::TPtr ApplyExtractMembersToMapNext(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
+    TCoMapNext mapNext(node);
+    YQL_CLOG(DEBUG, Core) << "Apply ExtractMembers to " << node->Content() << logSuffix;
+
+    const bool singleValue = GetItemType(*mapNext.Lambda().Ref().GetTypeAnn()) == nullptr;
+    TSet<TStringBuf> memberNames;
+    for (auto& member : members->ChildrenList()) {
+        YQL_ENSURE(member->IsAtom());
+        memberNames.insert(member->Content());
+    }
+    TExprBase newBody{FilterByFields(mapNext.Pos(), mapNext.Lambda().Body().Ptr(), memberNames, ctx, singleValue)};
+
+    return Build<TCoMapNext>(ctx, mapNext.Pos())
+        .Input(mapNext.Input())
+        .Lambda()
+            .Args({"current", "next"})
+            .Body<TExprApplier>()
+                .Apply(newBody)
+                .With(mapNext.Lambda().Args().Arg(0), "current")
+                .With(mapNext.Lambda().Args().Arg(1), "next")
+            .Build()
+        .Build()
+        .Done()
+        .Ptr();
+}
+
 } // NYql
diff --git a/ydb/library/yql/core/common_opt/yql_co_extr_members.h b/ydb/library/yql/core/common_opt/yql_co_extr_members.h
index 8244636784..974775b49d 100644
--- a/ydb/library/yql/core/common_opt/yql_co_extr_members.h
+++ b/ydb/library/yql/core/common_opt/yql_co_extr_members.h
@@ -25,5 +25,6 @@ TExprNode::TPtr ApplyExtractMembersToAggregate(const TExprNode::TPtr& node, cons
 TExprNode::TPtr ApplyExtractMembersToChopper(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
 TExprNode::TPtr ApplyExtractMembersToCollect(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
 TExprNode::TPtr ApplyExtractMembersToMapJoinCore(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
+TExprNode::TPtr ApplyExtractMembersToMapNext(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
 
 } // NYql
diff --git a/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp b/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp
index b77766bc6f..9a72235ea8 100644
--- a/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp
@@ -203,6 +203,14 @@ void RegisterCoFinalizers(TFinalizingOptimizerMap& map) {
             }
         );
     };
+
+    map[TCoMapNext::CallableName()] = [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
+        SubsetFieldsForNodeWithMultiUsage(node, *optCtx.ParentsMap, toOptimize, ctx,
+            [] (const TExprNode::TPtr& input, const TExprNode::TPtr& members, const TParentsMap&, TExprContext& ctx) {
+                return ApplyExtractMembersToMapNext(input, members, ctx, " with multi-usage");
+            }
+        );
+    };
 }
 
 } // NYql
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 88f9140031..690393bbfe 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
@@ -1937,6 +1937,13 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
             return node;
         }
 
+        if (self.Input().Maybe<TCoMapNext>()) {
+            if (auto res = ApplyExtractMembersToMapNext(self.Input().Ptr(), self.Members().Ptr(), ctx, {})) {
+                return res;
+            }
+            return node;
+        }
+
         return node;
     };
 
@@ -2302,6 +2309,37 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
         }
         return node;
     };
+
+    map[TCoMapNext::CallableName()] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        TCoMapNext self(node);
+        if (!optCtx.IsSingleUsage(self.Input().Ref())) {
+            return node;
+        }
+
+        std::map<std::string_view, TExprNode::TPtr> usedFields;
+        if ((
+             HaveFieldsSubset(self.Lambda().Body().Ptr(), self.Lambda().Args().Arg(0).Ref(), usedFields, *optCtx.ParentsMap, false) &&
+             HaveFieldsSubset(self.Lambda().Body().Ptr(), self.Lambda().Args().Arg(1).Ref(), usedFields, *optCtx.ParentsMap, false)
+            ) && usedFields.size() < GetSeqItemType(self.Input().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<TCoMapNext>(ctx, node->Pos())
+                .Input<TCoExtractMembers>()
+                    .Input(self.Input())
+                    .Members()
+                        .Add(std::move(fields))
+                    .Build()
+                .Build()
+                .Lambda(ctx.DeepCopyLambda(self.Lambda().Ref()))
+                .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 8a66d1bd60..11899138d8 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -7194,6 +7194,25 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
         return ctx.NewCallable(node->Pos(), "Member", { node->HeadPtr(), node->ChildPtr(1) });
     };
 
+    map["MapNext"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& /*optCtx*/) {
+        TCoMapNext self(node);
+        if (!IsDepended(self.Lambda().Body().Ref(), self.Lambda().Args().Arg(1).Ref())) {
+            YQL_CLOG(DEBUG, Core) << node->Content() << " with unused next arg";
+            return Build<TCoOrderedMap>(ctx, self.Pos())
+                .Input(self.Input())
+                .Lambda()
+                    .Args({"row"})
+                    .Body<TExprApplier>()
+                        .Apply(self.Lambda().Body())
+                        .With(self.Lambda().Args().Arg(0), "row")
+                    .Build()
+                .Build()
+                .Done()
+                .Ptr();
+        }
+        return node;
+    };
+
     // will be applied to any callable after all above
     map[""] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
         YQL_ENSURE(node->IsCallable());
diff --git a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
index edfcaa1530..96ed88efbb 100644
--- a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
+++ b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
@@ -183,6 +183,11 @@
             "Match": {"Type": "Callable", "Name": "OrderedMap"}
         },
         {
+            "Name": "TCoMapNext",
+            "Base": "TCoMapBase",
+            "Match": {"Type": "Callable", "Name": "MapNext"}
+        },
+        {
             "Name": "TCoMultiMap",
             "Base": "TCoMapBase",
             "Match": {"Type": "Callable", "Name": "MultiMap"}
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index 69aa84a500..8385c239ae 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -207,7 +207,8 @@ TExprNode::TPtr KeepColumnOrder(const TExprNode::TPtr& node, const TExprNode& sr
 
 template<class TFieldsSet>
 bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap, bool allowDependsOn) {
-    if (arg.GetTypeAnn()->GetKind() != ETypeAnnotationKind::Struct) {
+    const TTypeAnnotationNode* argType = RemoveOptionalType(arg.GetTypeAnn());
+    if (argType->GetKind() != ETypeAnnotationKind::Struct) {
         return false;
     }
 
@@ -215,7 +216,7 @@ bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TField
         return false;
     }
 
-    const auto inputStructType = arg.GetTypeAnn()->Cast<TStructExprType>();
+    const auto inputStructType = argType->Cast<TStructExprType>();
     if (!IsDepended(*start, arg)) {
         return inputStructType->GetSize() > 0;
     }