Support ExtractMembers in SqlIn optimizer. (KIKIMR-14325)

ref:8890678ec7a8099a91e01661b77847d10ccf1793

3 files changed, 58 insertions, 4 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 af9cf9fb9e1..802c6e6e732 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
@@ -350,6 +350,7 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
             .Build()
         .Done();
 
+    rightReadMatch->ExtractMembers = {}; // We already fetching only required columns
     lookup = rightReadMatch->BuildProcessNodes(lookup, ctx);
 
     if (join.JoinType().Value() == "RightSemi") {
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
index fd5d2130011..a09a982c562 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
@@ -1,3 +1,4 @@
+#include "kqp_opt_log_impl.h"
 #include "kqp_opt_log_rules.h"
 
 #include <ydb/core/kqp/opt/kqp_opt_impl.h>
@@ -40,19 +41,32 @@ TExprBase KqpRewriteSqlInToEquiJoin(const TExprBase& node, TExprContext& ctx, co
         return node;
     }
 
-    if (!flatMap.Input().Maybe<TKqlReadTable>() && !flatMap.Input().Maybe<TKqlReadTableIndex>()) {
+    auto readMatch = MatchRead<TKqlReadTableBase>(flatMap.Input());
+    if (!readMatch) {
         return node;
     }
 
-    const auto readTable = flatMap.Input().Cast<TKqlReadTableBase>();
+    if (readMatch->FlatMap) {
+        return node;
+    }
+
+    auto readTable = readMatch->Read.Cast<TKqlReadTableBase>();
+
+    static const std::set<TStringBuf> supportedReads {
+        TKqlReadTable::CallableName(),
+        TKqlReadTableIndex::CallableName(),
+    };
+
+    if (!supportedReads.contains(readTable.CallableName())) {
+        return node;
+    }
 
     if (!readTable.Table().SysView().Value().empty()) {
         return node;
     }
 
     TString lookupTable;
-
-    if (auto indexRead = flatMap.Input().Maybe<TKqlReadTableIndex>()) {
+    if (auto indexRead = readTable.Maybe<TKqlReadTableIndex>()) {
         lookupTable = GetIndexMetadata(indexRead.Cast(), *kqpCtx.Tables, kqpCtx.Cluster)->Name;
     } else {
         lookupTable = readTable.Table().Path().StringValue();
diff --git a/ydb/core/kqp/ut/kqp_indexes_ut.cpp b/ydb/core/kqp/ut/kqp_indexes_ut.cpp
index c2580fb2aaa..b33b433e426 100644
--- a/ydb/core/kqp/ut/kqp_indexes_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_indexes_ut.cpp
@@ -5226,6 +5226,45 @@ R"([[#;#;["Primary1"];[41u]];[["Secondary2"];[2u];["Primary2"];[42u]];[["Seconda
             [[2];[40u];["Two"];["Value4"]]
         ])", FormatResultSetYson(result.GetResultSet(0)));
     }
+
+    Y_UNIT_TEST_TWIN(IndexMultipleRead, UseNewEngine) {
+        TKikimrRunner kikimr;
+
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+        CreateSampleTablesWithIndex(session);
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$fks")
+                .BeginList()
+                .AddListItem().Int32(5)
+                .AddListItem().Int32(10)
+                .EndList()
+                .Build()
+            .Build();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            DECLARE $fks AS List<Int32>;
+
+            SELECT * FROM SecondaryKeys VIEW Index WHERE Fk IN $fks;
+            SELECT COUNT(*) FROM SecondaryKeys VIEW Index WHERE Fk IN $fks;
+        )"), TTxControl::BeginTx().CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        AssertTableStats(result, "/Root/SecondaryKeys", {
+            .ExpectedReads = UseNewEngine ? 2 : 1, // TODO: Looks like missing SkipNullMembers in NewEngine
+        });
+
+        AssertTableStats(result, "/Root/SecondaryKeys/Index/indexImplTable", {
+            .ExpectedReads = UseNewEngine ? 1 : 2,
+        });
+
+        CompareYson(R"([[[5];[5];["Payload5"]]])", FormatResultSetYson(result.GetResultSet(0)));
+        CompareYson(R"([[1u]])", FormatResultSetYson(result.GetResultSet(1)));
+    }
 }
 
 }