[yql] Fix CanEvaluate check in DQ/YT providers

YQL-12393

ref:f4b431413c42159144c6343b35ab9c302d27d96c

5 files changed, 58 insertions, 9 deletions

diff --git a/ydb/library/yql/core/services/yql_eval_expr.cpp b/ydb/library/yql/core/services/yql_eval_expr.cpp
index 2440e5ea25c..fce5d629feb 100644
--- a/ydb/library/yql/core/services/yql_eval_expr.cpp
+++ b/ydb/library/yql/core/services/yql_eval_expr.cpp
@@ -439,7 +439,7 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
                                 continue;
                             }
 
-                            if (!p->GetExecWorld(*x.first, calcWorldRoot)) {
+                            if (!p->GetExecWorld(x.first, calcWorldRoot)) {
                                 canExec = false;
                                 break;
                             }
@@ -472,12 +472,12 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
                 ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << "Only pure expressions are supported"));
                 return IGraphTransformer::TStatus::Error;
             }
+        }
 
-            if (!calcWorldRoot) {
-                calcWorldRoot = ctx.NewWorld(input->Pos());
-                calcWorldRoot->SetTypeAnn(ctx.MakeType<TUnitExprType>());
-                calcWorldRoot->SetState(TExprNode::EState::ConstrComplete);
-            }
+        if (!calcWorldRoot) {
+            calcWorldRoot = ctx.NewWorld(input->Pos());
+            calcWorldRoot->SetTypeAnn(ctx.MakeType<TUnitExprType>());
+            calcWorldRoot->SetState(TExprNode::EState::ConstrComplete);
         }
 
         return IGraphTransformer::TStatus::Ok;
diff --git a/ydb/library/yql/core/yql_data_provider.h b/ydb/library/yql/core/yql_data_provider.h
index 049e97c6388..f8ecac31f8a 100644
--- a/ydb/library/yql/core/yql_data_provider.h
+++ b/ydb/library/yql/core/yql_data_provider.h
@@ -140,7 +140,7 @@ public:
     // Right! or worlds are written to syncList
     virtual bool CanBuildResult(const TExprNode& node, TSyncMap& syncList) = 0;
     virtual bool CanPullResult(const TExprNode& node, TSyncMap& syncList, bool& canRef) = 0;
-    virtual bool GetExecWorld(const TExprNode& node, TExprNode::TPtr& root) = 0;
+    virtual bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) = 0;
     virtual bool CanEvaluate(const TExprNode& node) = 0;
     virtual void UndoEvaluationChanges() = 0;
     virtual TExprNode::TPtr CleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx) = 0;
diff --git a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
index 8de3d57fd3a..c78ca70e552 100644
--- a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
+++ b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
@@ -205,7 +205,7 @@ bool TDataProviderBase::CanPullResult(const TExprNode& node, TSyncMap& syncList,
     return false;
 }
 
-bool TDataProviderBase::GetExecWorld(const TExprNode& node, TExprNode::TPtr& root) {
+bool TDataProviderBase::GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) {
     root = nullptr;
     Y_UNUSED(node);
     return false;
diff --git a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
index 242a39877e1..5f30ced4500 100644
--- a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
+++ b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
@@ -64,7 +64,7 @@ public:
     bool IsPersistent(const TExprNode& node) override;
     bool CanBuildResult(const TExprNode& node, TSyncMap& syncList) override;
     bool CanPullResult(const TExprNode& node, TSyncMap& syncList, bool& canRef) override;
-    bool GetExecWorld(const TExprNode& node, TExprNode::TPtr& root) override;
+    bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) override;
     bool CanEvaluate(const TExprNode& node) override;
     void UndoEvaluationChanges() override;
     TExprNode::TPtr CleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx) override;
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_datasource.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_datasource.cpp
index aea55435de9..c1f905df2a8 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_datasource.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_datasource.cpp
@@ -58,6 +58,55 @@ public:
         return *ConfigurationTransformer;
     }
 
+    bool CanBuildResult(const TExprNode& node, TSyncMap& syncList) override {
+        if (!node.IsComplete()) {
+            return false;
+        }
+
+        bool canBuild = true;
+        VisitExpr(node, [&canBuild] (const TExprNode& n) {
+            if (!canBuild) {
+                return false;
+            }
+            if (TDqConnection::Match(&n) || TDqPhyPrecompute::Match(&n)) {
+                // Don't go deeper
+                return false;
+            }
+            if (!n.IsComposable()) {
+                canBuild = false;
+                return false;
+            }
+            return true;
+        });
+
+        if (canBuild) {
+            for (const auto& child : node.ChildrenList()) {
+                VisitExpr(child, [&syncList] (const TExprNode::TPtr& item) {
+                    if (ETypeAnnotationKind::World == item->GetTypeAnn()->GetKind()) {
+                        syncList.emplace(item, syncList.size());
+                        return false;
+                    }
+                    return true;
+                });
+            }
+        }
+
+        return canBuild;
+    }
+
+    bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) override {
+        if (ETypeAnnotationKind::World == node->GetTypeAnn()->GetKind() && !root) {
+            root = node;
+            return true;
+        }
+        root = nullptr;
+        return false;
+    }
+
+    bool CanEvaluate(const TExprNode& node) override {
+        return TDqConnection::Match(&node) || TDqPhyPrecompute::Match(&node);
+    }
+
     TExprNode::TPtr OptimizePull(const TExprNode::TPtr& node, const TFillSettings& fillSettings, TExprContext& ctx,
         IOptimizationContext& optCtx) override
     {