Fix missing ExtractMembers in KqpJoinToIndexLookup. ()

2 files changed, 42 insertions, 2 deletions

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
index 32bdb23e3c..f4f55af0cf 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
@@ -437,14 +437,18 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
             .Build()
         .Done();
 
-    rightReadMatch->ExtractMembers = {}; // We already fetching only required columns
+    auto lookupColumns = read.Columns();
+    if (rightReadMatch->ExtractMembers) {
+        lookupColumns = rightReadMatch->ExtractMembers.Cast().Members();
+    }
+
     lookup = rightReadMatch->BuildProcessNodes(lookup, ctx);
 
     if (join.JoinType().Value() == "RightSemi") {
         auto arg = TCoArgument(ctx.NewArgument(join.Pos(), "row"));
         auto rightLabel = join.RightLabel().Cast<TCoAtom>().Value();
 
-        TVector<TExprBase> renames = CreateRenames(rightReadMatch->FlatMap, read.Columns(), arg, rightLabel,
+        TVector<TExprBase> renames = CreateRenames(rightReadMatch->FlatMap, lookupColumns, arg, rightLabel,
             join.Pos(), ctx);
 
         lookup = Build<TCoMap>(ctx, join.Pos())
diff --git a/ydb/core/kqp/ut/kqp_ne_ut.cpp b/ydb/core/kqp/ut/kqp_ne_ut.cpp
index 1e5856790e..046d21b29c 100644
--- a/ydb/core/kqp/ut/kqp_ne_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_ut.cpp
@@ -3293,6 +3293,42 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
             UNIT_ASSERT_VALUES_EQUAL(FromString<i32>(read["limit"].GetString()), 10);
         }
     }
+
+    Y_UNIT_TEST(IdxLookupExtractMembers) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExplainDataQuery(R"(
+            --!syntax_v1
+            PRAGMA kikimr.UseNewEngine = 'true';
+
+            DECLARE $input AS List<Struct<
+                Key: Uint64,
+                Text: String,
+            >>;
+
+            $to_upsert = (
+                SELECT
+                    i.Key AS Key,
+                FROM AS_TABLE($input) AS i
+                JOIN EightShard AS s
+                USING (Key)
+                WHERE s.Key IS NULL OR s.Text != i.Text
+            );
+
+            $to_delete = (
+                SELECT s.Key AS Key
+                FROM EightShard AS s
+                JOIN AS_TABLE($input) AS i
+                USING (Key)
+            );
+
+            DELETE FROM EightShard ON SELECT Key FROM $to_delete;
+            UPSERT INTO EightShard SELECT * FROM $to_upsert;
+        )").ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+    }
 }
 
 } // namespace NKikimr::NKqp