More efficient DqBuildPrecompute optimizer rule. (KIKIMR-14484)

ref:1708b45a3689a4f93862059b6af9a7fec2f89044

11 files changed, 78 insertions, 51 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
index 1576f97467f..ee2f1fe6e6b 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
@@ -48,7 +48,7 @@ public:
         AddHandler(0, &TDqJoin::Match, HNDL(RewriteRightJoinToLeft));
         AddHandler(0, &TDqJoin::Match, HNDL(RewriteLeftPureJoin<false>));
         AddHandler(0, &TDqJoin::Match, HNDL(BuildJoin<false>));
-        AddHandler(0, &TDqPrecompute::Match, HNDL(BuildPrecomputeStage));
+        AddHandler(0, &TDqPrecompute::Match, HNDL(BuildPrecompute));
         AddHandler(0, &TCoLMap::Match, HNDL(PushLMapToStage<false>));
         AddHandler(0, &TCoOrderedLMap::Match, HNDL(PushOrderedLMapToStage<false>));
         AddHandler(0, &TKqlInsertRows::Match, HNDL(BuildInsertStages));
@@ -233,9 +233,9 @@ protected:
         return output;
     }
 
-    TMaybeNode<TExprBase> BuildPrecomputeStage(TExprBase node, TExprContext& ctx) {
-        TExprBase output = DqBuildPrecomputeStage(node, ctx);
-        DumpAppliedRule("BuildPrecomputeStage", node.Ptr(), output.Ptr(), ctx);
+    TMaybeNode<TExprBase> BuildPrecompute(TExprBase node, TExprContext& ctx) {
+        TExprBase output = DqBuildPrecompute(node, ctx);
+        DumpAppliedRule("BuildPrecompute", node.Ptr(), output.Ptr(), ctx);
         return output;
     }
 
diff --git a/ydb/core/kqp/ut/CMakeLists.darwin.txt b/ydb/core/kqp/ut/CMakeLists.darwin.txt
index 96c652a9f89..e06351bcca2 100644
--- a/ydb/core/kqp/ut/CMakeLists.darwin.txt
+++ b/ydb/core/kqp/ut/CMakeLists.darwin.txt
@@ -33,7 +33,6 @@ target_sources(ydb-core-kqp-ut PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_acl_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_arrow_in_channels_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_document_api_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_effects_perf_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_explain_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_flip_join_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_index_lookup_join_ut.cpp
@@ -44,10 +43,11 @@ target_sources(ydb-core-kqp-ut PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_locks_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_merge_connection_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_mvcc_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_inplace_update_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_effects_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_flowcontrol_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_effects_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_inplace_update_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_olap_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_params_ut.cpp
diff --git a/ydb/core/kqp/ut/CMakeLists.linux.txt b/ydb/core/kqp/ut/CMakeLists.linux.txt
index a7aa079935f..460c046d669 100644
--- a/ydb/core/kqp/ut/CMakeLists.linux.txt
+++ b/ydb/core/kqp/ut/CMakeLists.linux.txt
@@ -34,7 +34,6 @@ target_sources(ydb-core-kqp-ut PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_acl_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_arrow_in_channels_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_document_api_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_effects_perf_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_explain_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_flip_join_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_index_lookup_join_ut.cpp
@@ -45,10 +44,11 @@ target_sources(ydb-core-kqp-ut PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_locks_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_merge_connection_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_mvcc_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_inplace_update_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_effects_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_flowcontrol_ut.cpp
-  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_newengine_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_effects_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_inplace_update_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_ne_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_olap_ut.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/kqp_params_ut.cpp
diff --git a/ydb/core/kqp/ut/kqp_newengine_effects_ut.cpp b/ydb/core/kqp/ut/kqp_ne_effects_ut.cpp
index 17b315ac4d0..17b315ac4d0 100644
--- a/ydb/core/kqp/ut/kqp_newengine_effects_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_effects_ut.cpp
diff --git a/ydb/core/kqp/ut/kqp_newengine_flowcontrol_ut.cpp b/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
index 16c89eceb58..16c89eceb58 100644
--- a/ydb/core/kqp/ut/kqp_newengine_flowcontrol_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
diff --git a/ydb/core/kqp/ut/kqp_newengine_inplace_update_ut.cpp b/ydb/core/kqp/ut/kqp_ne_inplace_update_ut.cpp
index 79a8839c596..632221d5a2b 100644
--- a/ydb/core/kqp/ut/kqp_newengine_inplace_update_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_inplace_update_ut.cpp
@@ -418,4 +418,4 @@ Y_UNIT_TEST_TWIN(BigRow, EnableInplaceUpdate) {
 } // suite
 
 } // namespace NKqp
-} // namespace NKikimr
\ No newline at end of file
+} // namespace NKikimr
diff --git a/ydb/core/kqp/ut/kqp_effects_perf_ut.cpp b/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
index 1b2691d1203..9815e322a9d 100644
--- a/ydb/core/kqp/ut/kqp_effects_perf_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
@@ -119,7 +119,7 @@ TParams BuildInsertIndexParams(TTableClient& client) {
 
 } // namespace
 
-Y_UNIT_TEST_SUITE(KqpEffectsPerf) {
+Y_UNIT_TEST_SUITE(KqpPerf) {
     Y_UNIT_TEST_TWIN(Upsert, UseNewEngine) {
         TKikimrRunner kikimr;
         auto db = kikimr.GetTableClient();
@@ -470,6 +470,59 @@ Y_UNIT_TEST_SUITE(KqpEffectsPerf) {
         auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
         UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), UseNewEngine ? 5 : 3);
     }
+
+    Y_UNIT_TEST_TWIN(IdxLookupJoin, UseNewEngine) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$key").Int32(3).Build()
+            .Build();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            DECLARE $key AS Int32;
+
+            SELECT *
+            FROM Join1 AS t1
+            INNER JOIN Join2 AS t2 ON t1.Fk21 = t2.Key1 AND t1.Fk22 = t2.Key2
+            WHERE t1.Key = $key;
+        )"), TTxControl::BeginTx().CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+        UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), UseNewEngine ? 3 : 3);
+    }
+
+    Y_UNIT_TEST_TWIN(IdxLookupJoinThreeWay, UseNewEngine) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$key").Int32(3).Build()
+            .Build();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            DECLARE $key AS Int32;
+
+            SELECT t1.Key, t3.Value
+            FROM Join1 AS t1
+            INNER JOIN Join2 AS t2 ON t1.Fk21 = t2.Key1 AND t1.Fk22 = t2.Key2
+            INNER JOIN KeyValue2 AS t3 ON t2.Name = t3.Key
+            WHERE t1.Key = $key;
+        )"), TTxControl::BeginTx().CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+        UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), UseNewEngine ? 5 : 4);
+    }
 }
 
 } // namespace NKikimr::NKqp
diff --git a/ydb/core/kqp/ut/kqp_newengine_ut.cpp b/ydb/core/kqp/ut/kqp_ne_ut.cpp
index 1043baef5e8..1043baef5e8 100644
--- a/ydb/core/kqp/ut/kqp_newengine_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_ut.cpp
diff --git a/ydb/core/kqp/ut/ya.make b/ydb/core/kqp/ut/ya.make
index 5aa9abf952a..7ada0c7f44a 100644
--- a/ydb/core/kqp/ut/ya.make
+++ b/ydb/core/kqp/ut/ya.make
@@ -22,7 +22,6 @@ SRCS(
     kqp_acl_ut.cpp
     kqp_arrow_in_channels_ut.cpp
     kqp_document_api_ut.cpp
-    kqp_effects_perf_ut.cpp
     kqp_explain_ut.cpp
     kqp_flip_join_ut.cpp
     kqp_index_lookup_join_ut.cpp
@@ -33,10 +32,11 @@ SRCS(
     kqp_locks_ut.cpp
     kqp_merge_connection_ut.cpp
     kqp_mvcc_ut.cpp
-    kqp_newengine_inplace_update_ut.cpp
-    kqp_newengine_effects_ut.cpp
-    kqp_newengine_flowcontrol_ut.cpp
-    kqp_newengine_ut.cpp
+    kqp_ne_effects_ut.cpp
+    kqp_ne_flowcontrol_ut.cpp
+    kqp_ne_inplace_update_ut.cpp
+    kqp_ne_perf_ut.cpp
+    kqp_ne_ut.cpp
     kqp_not_null_columns_ut.cpp
     kqp_olap_ut.cpp
     kqp_params_ut.cpp
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index ed39039ed3d..015a36c3004 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -1424,9 +1424,9 @@ TExprBase DqBuildExtendStage(TExprBase node, TExprContext& ctx) {
 }
 
 /*
- * Precompute input value in a separate stage.
+ * Build physical precompute node.
  */
-TExprBase DqBuildPrecomputeStage(TExprBase node, TExprContext& ctx) {
+TExprBase DqBuildPrecompute(TExprBase node, TExprContext& ctx) {
     if (!node.Maybe<TDqPrecompute>()) {
         return node;
     }
@@ -1434,13 +1434,10 @@ TExprBase DqBuildPrecomputeStage(TExprBase node, TExprContext& ctx) {
     auto input = node.Cast<TDqPrecompute>().Input();
 
     TExprNode::TPtr connection;
-    bool value = false;
-
     if (input.Maybe<TDqCnUnionAll>()) {
         connection = input.Ptr();
     } else if (input.Maybe<TDqCnValue>()) {
         connection = input.Ptr();
-        value = true;
     } else if (IsDqPureExpr(input)) {
         if (input.Ref().GetTypeAnn()->GetKind() != ETypeAnnotationKind::List &&
             input.Ref().GetTypeAnn()->GetKind() != ETypeAnnotationKind::Data)
@@ -1469,7 +1466,6 @@ TExprBase DqBuildPrecomputeStage(TExprBase node, TExprContext& ctx) {
                 .Build()
             .Done().Ptr();
 
-        value = true;
     } else {
         return node;
     }
@@ -1478,29 +1474,7 @@ TExprBase DqBuildPrecomputeStage(TExprBase node, TExprContext& ctx) {
         .Connection(connection)
         .Done();
 
-    if (value) {
-        return phyPrecompute;
-    }
-
-    auto precomputeStage = Build<TDqStage>(ctx, node.Pos())
-        .Inputs()
-            .Add(phyPrecompute)
-            .Build()
-        .Program()
-            .Args({"zzz"})
-            .Body<TCoIterator>()
-                .List("zzz")
-                .Build()
-            .Build()
-        .Settings().Build()
-        .Done();
-
-    return Build<TDqCnUnionAll>(ctx, node.Pos())
-        .Output()
-            .Stage(precomputeStage)
-            .Index().Build("0")
-            .Build()
-        .Done();
+    return phyPrecompute;
 }
 
 TExprBase DqBuildHasItems(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx) {
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.h b/ydb/library/yql/dq/opt/dq_opt_phy.h
index 550410253d7..4a9ac46701e 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.h
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.h
@@ -79,7 +79,7 @@ NNodes::TExprBase DqBuildPureExprStage(NNodes::TExprBase node, TExprContext& ctx
 
 NNodes::TExprBase DqBuildExtendStage(NNodes::TExprBase node, TExprContext& ctx);
 
-NNodes::TExprBase DqBuildPrecomputeStage(NNodes::TExprBase node, TExprContext& ctx);
+NNodes::TExprBase DqBuildPrecompute(NNodes::TExprBase node, TExprContext& ctx);
 
 NYql::NNodes::TExprBase DqBuildHasItems(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx,
     NYql::IOptimizationContext& optCtx);