YQL-19476: Fix flatten columns with udf

commit_hash:7f6bda02fbc849702796300dfa7f1e78b361bd79

1 files changed, 5 insertions, 3 deletions

diff --git a/yql/essentials/core/services/yql_lineage.cpp b/yql/essentials/core/services/yql_lineage.cpp
index c78bbb6b820..12f624c5c17 100644
--- a/yql/essentials/core/services/yql_lineage.cpp
+++ b/yql/essentials/core/services/yql_lineage.cpp
@@ -543,9 +543,11 @@ private:
             if (lambda.GetTypeAnn()->GetKind() == ETypeAnnotationKind::List) {
                 value = &body;
                 while(value->IsCallable({"FlatMap", "OrderedFlatMap"})) {
-                    if (value->Head().IsCallable("Member") && &value->Head().Head() == &arg) {
-                        TString field(value->Head().Tail().Content());
-                        flattenColumns.emplace(value->Tail().Head().HeadPtr().Get(), field);
+                    TNodeMap<TMaybe<TFieldsLineage>> visited;
+                    if (auto res = ScanExprLineage(value->Head(), &arg, &innerLineage, visited, {})) {
+                        for (const auto& f: res->Items) {
+                            flattenColumns.emplace(value->Tail().Head().HeadPtr().Get(), f.Field);
+                        }
                     }
                     value = &value->Tail().Tail();
                 }