Apply ExtractMembers to NarrowMap.

3 files changed, 32 insertions, 0 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 67d9d9e1c8..afbbf1a586 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
@@ -947,4 +947,14 @@ TExprNode::TPtr ApplyExtractMembersToCombineCore(const TExprNode::TPtr& node, co
         .Ptr();
 }
 
+TExprNode::TPtr ApplyExtractMembersToNarrowMap(const TExprNode::TPtr& node, const TExprNode::TPtr& members, bool isFlat, TExprContext& ctx, TStringBuf logSuffix) {
+    YQL_CLOG(DEBUG, Core) << "Apply ExtractMembers to " << node->Content() << logSuffix;
+    auto body = GetLambdaBody(node->Tail());
+    std::for_each(body.begin(), body.end(), [&](TExprNode::TPtr& root) {
+        const auto pos = root->Pos();
+        root = ctx.NewCallable(pos, isFlat ? TCoExtractMembers::CallableName() : TCoFilterMembers::CallableName(), {std::move(root), members});
+    });
+    return ctx.ChangeChild(*node, TCoMapBase::idx_Lambda, ctx.DeepCopyLambda(node->Tail(), std::move(body)));
+}
+
 } // 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 91b762587c..530979f65d 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
@@ -29,5 +29,6 @@ TExprNode::TPtr ApplyExtractMembersToMapNext(const TExprNode::TPtr& node, const
 TExprNode::TPtr ApplyExtractMembersToChain1Map(const TExprNode::TPtr& node, TExprNode::TPtr members, const TParentsMap& parentsMap, TExprContext& ctx, TStringBuf logSuffix);
 TExprNode::TPtr ApplyExtractMembersToCondense1(const TExprNode::TPtr& node, TExprNode::TPtr members, const TParentsMap& parentsMap, TExprContext& ctx, TStringBuf logSuffix);
 TExprNode::TPtr ApplyExtractMembersToCombineCore(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
+TExprNode::TPtr ApplyExtractMembersToNarrowMap(const TExprNode::TPtr& node, const TExprNode::TPtr& members, bool isFlat, TExprContext& ctx, TStringBuf logSuffix);
 
 } // 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 c3cb4f6b04..1b41f326c0 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
@@ -1494,6 +1494,27 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
             return node;
         }
 
+        if (self.Input().Maybe<TCoNarrowMap>()) {
+            if (auto res = ApplyExtractMembersToNarrowMap(self.Input().Ptr(), self.Members().Ptr(), false, ctx, {})) {
+                return res;
+            }
+            return node;
+        }
+
+        if (self.Input().Maybe<TCoNarrowMultiMap>()) {
+            if (auto res = ApplyExtractMembersToNarrowMap(self.Input().Ptr(), self.Members().Ptr(), false, ctx, {})) {
+                return res;
+            }
+            return node;
+        }
+
+        if (const auto narrow = self.Input().Maybe<TCoNarrowFlatMap>()) {
+            if (auto res = ApplyExtractMembersToNarrowMap(self.Input().Ptr(), self.Members().Ptr(), ETypeAnnotationKind::Optional != narrow.Cast().Lambda().Body().Ref().GetTypeAnn()->GetKind(), ctx, {})) {
+                return res;
+            }
+            return node;
+        }
+
         return node;
     };