Add KqpApplyExtractMembersToLookupTable optimizer rule. (KIKIMR-14298)

ref:e5ea3df9ffd85bde8a912a2c3ba6b836ea249faf

4 files changed, 67 insertions, 0 deletions

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
index d951c5ce30..793b8a9343 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -31,6 +31,7 @@ public:
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ApplyExtractMembersToReadTable));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ApplyExtractMembersToReadTableRanges));
         AddHandler(0, &TCoExtractMembers::Match, HNDL(ApplyExtractMembersToReadOlapTable));
+        AddHandler(0, &TCoExtractMembers::Match, HNDL(ApplyExtractMembersToLookupTable));
         AddHandler(0, &TCoTake::Match, HNDL(RewriteTakeSortToTopSort));
         AddHandler(0, &TCoFlatMap::Match, HNDL(RewriteSqlInToEquiJoin));
         AddHandler(0, &TCoFlatMap::Match, HNDL(RewriteSqlInCompactToJoin));
@@ -87,6 +88,12 @@ protected:
         return output;
     }
 
+    TMaybeNode<TExprBase> ApplyExtractMembersToLookupTable(TExprBase node, TExprContext& ctx) {
+        TExprBase output = KqpApplyExtractMembersToLookupTable(node, ctx);
+        DumpAppliedRule("ApplyExtractMembersToLookupTable", node.Ptr(), output.Ptr(), ctx);
+        return output;
+    }
+
     TMaybeNode<TExprBase> RewriteTakeSortToTopSort(TExprBase node, TExprContext& ctx, const TGetParents& getParents) {
         TExprBase output = DqRewriteTakeSortToTopSort(node, ctx, *getParents());
         DumpAppliedRule("RewriteTakeSortToTopSort", node.Ptr(), output.Ptr(), ctx);
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp
index 96642998db..fcfd8a9e92 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp
@@ -111,5 +111,38 @@ TExprBase KqpApplyExtractMembersToReadOlapTable(TExprBase node, TExprContext& ct
         .Done();
 }
 
+TExprBase KqpApplyExtractMembersToLookupTable(TExprBase node, TExprContext& ctx) {
+    if (!node.Maybe<TCoExtractMembers>()) {
+        return node;
+    }
+
+    auto extract = node.Cast<TCoExtractMembers>();
+    auto input = extract.Input();
+
+    if (!input.Maybe<TKqlLookupTableBase>()) {
+        return node;
+    }
+
+    auto lookup = extract.Input().Cast<TKqlLookupTableBase>();
+
+    if (auto maybeIndexLookup = lookup.Maybe<TKqlLookupIndex>()) {
+        auto indexLookup = maybeIndexLookup.Cast();
+
+        return Build<TKqlLookupIndex>(ctx, lookup.Pos())
+            .Table(indexLookup.Table())
+            .LookupKeys(indexLookup.LookupKeys())
+            .Columns(extract.Members())
+            .Index(indexLookup.Index())
+            .Done();
+    }
+
+    return Build<TKqlLookupTableBase>(ctx, lookup.Pos())
+        .CallableName(lookup.CallableName())
+        .Table(lookup.Table())
+        .LookupKeys(lookup.LookupKeys())
+        .Columns(extract.Members())
+        .Done();
+}
+
 } // namespace NKikimr::NKqp::NOpt
 
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h b/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
index 9e84d83d82..75420f5cfe 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
@@ -23,6 +23,8 @@ NYql::NNodes::TExprBase KqpApplyExtractMembersToReadOlapTable(NYql::NNodes::TExp
 
 NYql::NNodes::TExprBase KqpApplyExtractMembersToReadTableRanges(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx);
 
+NYql::NNodes::TExprBase KqpApplyExtractMembersToLookupTable(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx);
+
 NYql::NNodes::TExprBase KqpJoinToIndexLookup(const NYql::NNodes::TExprBase& node, NYql::TExprContext& ctx,
     const TKqpOptimizeContext& kqpCtx, const NYql::TKikimrConfiguration::TPtr& config);
 
diff --git a/ydb/core/kqp/ut/kqp_newengine_ut.cpp b/ydb/core/kqp/ut/kqp_newengine_ut.cpp
index 6a239188d5..9e2e136ccc 100644
--- a/ydb/core/kqp/ut/kqp_newengine_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_newengine_ut.cpp
@@ -3019,6 +3019,31 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
         CompareYson(R"([[[1u];["One"]];[[2u];["Two"]]])", FormatResultSetYson(result.GetResultSet(3)));
         CompareYson(R"([[[3500u];["None"];[1u];["Anna"]];[[300u];["None"];[1u];["Paul"]];[[7200u];["None"];[2u];["Tony"]]])", FormatResultSetYson(result.GetResultSet(4)));
     }
+
+    Y_UNIT_TEST(LookupColumns) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExplainDataQuery(R"(
+            --!syntax_v1
+            PRAGMA kikimr.UseNewEngine = 'true';
+
+            $current_value = SELECT Fk21 FROM Join1 WHERE Key = 2;
+
+            UPDATE Join1 SET Fk22 = "New" WHERE Key = 1 AND Fk21 = 100;
+
+            SELECT $current_value;
+        )").ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        NJson::TJsonValue plan;
+        NJson::ReadJsonTree(result.GetPlan(), &plan, true);
+        auto reads = plan["tables"][0]["reads"].GetArraySafe();
+        for (auto& read : reads) {
+            UNIT_ASSERT(read["columns"].GetArraySafe().size() <= 2);
+        }
+    }
 }
 
 } // namespace NKikimr::NKqp