QLFilter: Sync HasNodesToCalculate with GetNodesToCalculate

commit_hash:74d863d4d6b21422b18a02c3b8315765448c030b

1 files changed, 44 insertions, 9 deletions

diff --git a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
index 2da17201ab4..aa9f5999538 100644
--- a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
@@ -388,14 +388,23 @@ TMaybe<TString> DeriveClusterFromSection(const NNodes::TYtSection& section, ERun
     return result;
 }
 
+bool IsQLFilterCompatibleOperation(const TExprNode& node) {
+    return node.IsCallable({"And", "Or", "Not", "Coalesce", "Exists", "<", "<=", ">", ">=", "==", "!="});
+}
+
+bool IsQLFilterRowMember(const TExprNode& node, const TExprNode* rowArg) {
+    return node.IsCallable("Member") && node.Child(0) == rowArg;
+}
+
 void GetNodesToCalculateFromQLFilter(const TExprNode& qlFilter, TExprNode::TListType &needCalc, TNodeSet &uniqNodes) {
     YQL_ENSURE(qlFilter.IsCallable("YtQLFilter"));
+    const auto rowArg = qlFilter.Child(1)->Child(0)->Child(0);
     const auto lambdaBody = qlFilter.Child(1)->Child(1);
-    VisitExpr(lambdaBody, [&needCalc, &uniqNodes](const TExprNode::TPtr& node) {
-        if (node->IsCallable({"And", "Or", "Not", "Coalesce", "Exists", "<", "<=", ">", ">=", "==", "!="})) {
+    VisitExpr(lambdaBody, [&needCalc, &uniqNodes, rowArg](const TExprNode::TPtr& node) {
+        if (IsQLFilterCompatibleOperation(*node)) {
             return true;
         }
-        if (node->IsCallable("Member")) {
+        if (IsQLFilterRowMember(*node, rowArg)) {
             return false;
         }
         if (uniqNodes.insert(node.Get()).second) {
@@ -407,6 +416,29 @@ void GetNodesToCalculateFromQLFilter(const TExprNode& qlFilter, TExprNode::TList
     });
 }
 
+bool HasNodesToCalculateFromQLFilter(const TExprNode& qlFilter) {
+    YQL_ENSURE(qlFilter.IsCallable("YtQLFilter"));
+    bool needCalc = false;
+    const auto rowArg = qlFilter.Child(1)->Child(0)->Child(0);
+    const auto lambdaBody = qlFilter.Child(1)->Child(1);
+    VisitExpr(lambdaBody, [&needCalc, rowArg](const TExprNode::TPtr& node) {
+        if (needCalc) {
+            return false;
+        }
+        if (IsQLFilterCompatibleOperation(*node)) {
+            return true;
+        }
+        if (IsQLFilterRowMember(*node, rowArg)) {
+            return false;
+        }
+        if (NeedCalc(TExprBase(node.Get()))) {
+            needCalc = true;
+        }
+        return false;
+    });
+    return needCalc;
+}
+
 } // unnamed
 
 TString GetClusterFromSection(const NNodes::TYtSection& section) {
@@ -798,12 +830,9 @@ TExprNode::TListType GetNodesToCalculate(const TExprNode::TPtr& input) {
                     break;
                 }
             }
-            for (const auto& path: section.Paths()) {
-                const TYtPathInfo pathInfo(path);
-                if (pathInfo.QLFilter) {
-                    GetNodesToCalculateFromQLFilter(*pathInfo.QLFilter, needCalc, uniqNodes);
-                }
-            }
+        }
+        else if (auto maybeQLFilter = TMaybeNode<TYtQLFilter>(node)) {
+            GetNodesToCalculateFromQLFilter(maybeQLFilter.Ref(), needCalc, uniqNodes);
         }
         else if (TMaybeNode<TYtOutput>(node)) {
             // Stop traversing dependent operations
@@ -876,6 +905,12 @@ bool HasNodesToCalculate(const TExprNode::TPtr& input) {
                 }
             }
         }
+        else if (auto maybeQLFilter = TMaybeNode<TYtQLFilter>(node)) {
+            if (HasNodesToCalculateFromQLFilter(maybeQLFilter.Ref())) {
+                needCalc = true;
+                return false;
+            }
+        }
         else if (TMaybeNode<TYtOutput>(node)) {
             // Stop traversing dependent operations
             return false;