Move TrimReadWorld rule to physical optimizers to not conflict with KQP optimizers

6 files changed, 37 insertions, 32 deletions

diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp
index c665badb40e..8650ae9b2f8 100644
--- a/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp
+++ b/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp
@@ -121,21 +121,11 @@ public:
         , State_(state)
     {
 #define HNDL(name) "LogicalOptimizer-"#name, Hndl(&TPqLogicalOptProposalTransformer::name)
-        AddHandler(0, &TCoLeft::Match, HNDL(TrimReadWorld));
       //  AddHandler(0, &TCoExtractMembers::Match, HNDL(ExtractMembers));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ExtractMembersOverDqWrap));
         #undef HNDL
     }
 
-    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
-        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TPqReadTopic>();
-        if (!maybeRead) {
-            return node;
-        }
-
-        return TExprBase(ctx.NewWorld(node.Pos()));
-    }
-
     /*
     TMaybeNode<TExprBase> ExtractMembers(TExprBase node, TExprContext& ctx) const {
         const auto& extract = node.Cast<TCoExtractMembers>();
diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_physical_optimize.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_physical_optimize.cpp
index 2f57e3999de..9205b65de9c 100644
--- a/ydb/library/yql/providers/pq/provider/yql_pq_physical_optimize.cpp
+++ b/ydb/library/yql/providers/pq/provider/yql_pq_physical_optimize.cpp
@@ -26,12 +26,24 @@ public:
         , State_(std::move(state))
     {
 #define HNDL(name) "PhysicalOptimizer-"#name, Hndl(&TPqPhysicalOptProposalTransformer::name)
+        AddHandler(0, &TCoLeft::Match, HNDL(TrimReadWorld));
         AddHandler(0, &TPqWriteTopic::Match, HNDL(PqWriteTopic));
 #undef HNDL
 
         SetGlobal(0); // Stage 0 of this optimizer is global => we can remap nodes.
     }
 
+    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
+        Y_UNUSED(ctx);
+
+        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TPqReadTopic>();
+        if (!maybeRead) {
+            return node;
+        }
+
+        return TExprBase(maybeRead.Cast().World().Ptr());
+    }
+
     NNodes::TCoNameValueTupleList BuildTopicWriteSettings(const TString& cluster, TPositionHandle pos, TExprContext& ctx) const {
         TVector<TCoNameValueTuple> props;
 
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
index bd3e92d89a4..af12022dea0 100644
--- a/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
+++ b/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
@@ -165,7 +165,6 @@ public:
         , State_(state)
     {
 #define HNDL(name) "LogicalOptimizer-"#name, Hndl(&TS3LogicalOptProposalTransformer::name)
-        AddHandler(0, &TCoLeft::Match, HNDL(TrimReadWorld));
         AddHandler(0, &TCoFlatMapBase::Match, HNDL(TryPrunePaths));
         AddHandler(0, &TDqSourceWrap::Match, HNDL(ApplyPrunedPath));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ExtractMembersOverDqSource));
@@ -282,15 +281,6 @@ public:
         return TStatus::Ok;
     }
 
-    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
-        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TS3ReadObject>();
-        if (!maybeRead) {
-            return node;
-        }
-
-        return TExprBase(ctx.NewWorld(node.Pos()));
-    }
-
     TMaybeNode<TExprBase> ApplyPrunedPath(TExprBase node, TExprContext& ctx) const {
         const TDqSourceWrap dqSource = node.Cast<TDqSourceWrap>();
         if (dqSource.DataSource().Category() != S3ProviderName) {
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp
index 7389c5b55b7..7bfad7c4e6c 100644
--- a/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp
+++ b/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp
@@ -102,10 +102,22 @@ public:
         , State_(std::move(state))
     {
 #define HNDL(name) "PhysicalOptimizer-"#name, Hndl(&TS3PhysicalOptProposalTransformer::name)
+        AddHandler(0, &TCoLeft::Match, HNDL(TrimReadWorld));
         AddHandler(0, &TS3WriteObject::Match, HNDL(S3WriteObject));
 #undef HNDL
     }
 
+    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
+        Y_UNUSED(ctx);
+
+        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TS3ReadObject>();
+        if (!maybeRead) {
+            return node;
+        }
+
+        return TExprBase(maybeRead.Cast().World().Ptr());
+    }
+
     TMaybeNode<TExprBase> S3WriteObject(TExprBase node, TExprContext& ctx, const TGetParents& getParents) const {
         const auto& write = node.Cast<TS3WriteObject>();
         const auto& targetNode = write.Target();
diff --git a/ydb/library/yql/providers/ydb/provider/yql_ydb_logical_opt.cpp b/ydb/library/yql/providers/ydb/provider/yql_ydb_logical_opt.cpp
index 512c126a751..b9c18d1ec11 100644
--- a/ydb/library/yql/providers/ydb/provider/yql_ydb_logical_opt.cpp
+++ b/ydb/library/yql/providers/ydb/provider/yql_ydb_logical_opt.cpp
@@ -24,21 +24,11 @@ public:
         , State_(state)
     {
 #define HNDL(name) "LogicalOptimizer-"#name, Hndl(&TYdbLogicalOptProposalTransformer::name)
-        AddHandler(0, &TCoLeft::Match, HNDL(TrimReadWorld));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ExtractMembers));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ExtractMembersOverDqWrap));
 #undef HNDL
     }
 
-    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
-        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TYdbReadTable>();
-        if (!maybeRead) {
-            return node;
-        }
-
-        return TExprBase(ctx.NewWorld(node.Pos()));
-    }
-
     TMaybeNode<TExprBase> ExtractMembers(TExprBase node, TExprContext& ctx) const {
         const auto& extract = node.Cast<TCoExtractMembers>();
         const auto& read = extract.Input().Maybe<TCoRight>().Input().Maybe<TYdbReadTable>();
diff --git a/ydb/library/yql/providers/ydb/provider/yql_ydb_physical_opt.cpp b/ydb/library/yql/providers/ydb/provider/yql_ydb_physical_opt.cpp
index a59751a8e05..3665754e793 100644
--- a/ydb/library/yql/providers/ydb/provider/yql_ydb_physical_opt.cpp
+++ b/ydb/library/yql/providers/ydb/provider/yql_ydb_physical_opt.cpp
@@ -25,12 +25,24 @@ public:
         , State_(state)
     {
 #define HNDL(name) "PhysicalOptimizer-"#name, Hndl(&TYdbPhysicalOptProposalTransformer::name)
+        AddHandler(0, &TCoLeft::Match,      HNDL(TrimReadWorld));
         AddHandler(0, &TCoTake::Match,      HNDL(Take));
         AddHandler(0, &TCoNarrowMap::Match, HNDL(ReadZeroColumns));
-        AddHandler(0, &TDqStage::Match, HNDL(SourceZeroColumns));
+        AddHandler(0, &TDqStage::Match,     HNDL(SourceZeroColumns));
 #undef HNDL
     }
 
+    TMaybeNode<TExprBase> TrimReadWorld(TExprBase node, TExprContext& ctx) const {
+        Y_UNUSED(ctx);
+
+        const auto& maybeRead = node.Cast<TCoLeft>().Input().Maybe<TYdbReadTable>();
+        if (!maybeRead) {
+            return node;
+        }
+
+        return TExprBase(maybeRead.Cast().World().Ptr());
+    }
+
     TMaybeNode<TExprBase> Take(TExprBase node, TExprContext& ctx) const {
         const auto& take = node.Cast<TCoTake>();
         const auto& wrap = take.Input().Maybe<TDqReadWrap>();
@@ -146,4 +158,3 @@ THolder<IGraphTransformer> CreateYdbPhysicalOptProposalTransformer(TYdbState::TP
 }
 
 } // namespace NYql
-