KIKIMR-14630: get rid of separate phase for parameters repack for sql-in queries

ref:1e6d6e8c02a300c29c4dc70197bd6216a42816d8

2 files changed, 111 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 802c6e6e73..bc0241286a 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
@@ -141,6 +141,36 @@ TVector<TExprBase> CreateRenames(const TMaybeNode<TCoFlatMap>& rightFlatmap, con
     return renames;
 }
 
+bool IsParameterToListOfStructsRepack(const TExprBase& expr) {
+    // Looking for next patterns:
+    //  - (FlatMap $in (lambda '($x) (Just (AsStruct '('"key" $x)))))
+    //  - (FlatMap $in (lambda '($x) (Just (AsStruct '('"key" (Nth $x '0)) '('"key2" (Nth $x '1))) ...)))
+
+    if (!expr.Maybe<TCoFlatMap>().Input().Maybe<TCoParameter>()) {
+        return false;
+    }
+    auto lambda = expr.Cast<TCoFlatMap>().Lambda();
+    if (lambda.Args().Size() != 1) {
+        return false;
+    }
+    if (!lambda.Body().Maybe<TCoJust>().Input().Maybe<TCoAsStruct>()) {
+        return false;
+    }
+    auto lambdaArg = lambda.Args().Arg(0).Raw();
+    auto asStruct = lambda.Body().Cast<TCoJust>().Input().Cast<TCoAsStruct>();
+
+    for (const auto& member : asStruct.Args()) {
+        if (member->Child(1) == lambdaArg) {
+            continue;
+        }
+        if (member->Child(1)->IsCallable("Nth") && member->Child(1)->Child(0) == lambdaArg) {
+            continue;
+        }
+        return false;
+    }
+
+    return true;
+}
 
 //#define DBG(...) YQL_CLOG(DEBUG, ProviderKqp) << __VA_ARGS__
 #define DBG(...)
@@ -281,12 +311,18 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
         return {};
     }
 
-    auto leftData = Build<TDqPrecompute>(ctx, join.Pos())
-        .Input(join.LeftInput())
-        .Done();
+    bool needPrecomputeLeft = !join.LeftInput().Maybe<TCoParameter>() &&
+                              !IsParameterToListOfStructsRepack(join.LeftInput());
+
+    TExprBase leftData = needPrecomputeLeft
+        ? Build<TDqPrecompute>(ctx, join.Pos())
+            .Input(join.LeftInput())
+            .Done()
+        : join.LeftInput();
+
     auto leftDataDeduplicated = DeduplicateByMembers(leftData, leftJoinKeys, ctx, join.Pos());
 
-    if (!equalLeftKeys.empty())    {
+    if (!equalLeftKeys.empty()) {
         auto row = Build<TCoArgument>(ctx, join.Pos())
             .Name("row")
             .Done();
diff --git a/ydb/core/kqp/ut/kqp_sqlin_ut.cpp b/ydb/core/kqp/ut/kqp_sqlin_ut.cpp
index b41df7c39e..1cb635dc7c 100644
--- a/ydb/core/kqp/ut/kqp_sqlin_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_sqlin_ut.cpp
@@ -1009,6 +1009,77 @@ Y_UNIT_TEST_SUITE(KqpSqlIn) {
             .ExpectedReads = UseNewEngine ? 1 : 6,
         });
     }
+
+    Y_UNIT_TEST_NEW_ENGINE(PhasesCount) {
+        TKikimrRunner kikimr;
+        auto session = kikimr.GetTableClient().CreateSession().GetValueSync().GetSession();
+
+        // simple key
+        {
+            const auto query = Q1_(R"(
+                DECLARE $in AS List<Uint64>;
+                SELECT * FROM `/Root/KeyValue` WHERE Key IN $in
+            )");
+
+            const auto params = TParamsBuilder()
+                .AddParam("$in")
+                    .BeginList()
+                        .AddListItem().Uint64(1)
+                        .AddListItem().Uint64(1)
+                        .AddListItem().Uint64(100)
+                    .EndList().Build()
+                .Build();
+
+            const auto settings = TExecDataQuerySettings()
+                .CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), params, settings).GetValueSync();
+
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            CompareYson(R"([[[1u];["One"]]])", FormatResultSetYson(result.GetResultSet(0)));
+
+            const Ydb::TableStats::QueryStats stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_EQUAL_C(2, stats.query_phases_size(), stats.DebugString());
+        }
+
+        // complex (tuple) key
+        {
+            const auto query = Q1_(R"(
+                DECLARE $in AS List<Tuple<UInt32, String>>;
+
+                SELECT * FROM `/Root/Test` WHERE (Group, Name) IN $in
+            )");
+
+            const auto params = TParamsBuilder()
+                .AddParam("$in")
+                    .BeginList()
+                        .AddListItem().BeginTuple()
+                            .AddElement().Uint32(1)
+                            .AddElement().String("Anna")
+                        .EndTuple()
+                        .AddListItem().BeginTuple()
+                            .AddElement().Uint32(1)
+                            .AddElement().String("Anna")
+                        .EndTuple()
+                        .AddListItem().BeginTuple()
+                            .AddElement().Uint32(2)
+                            .AddElement().String("Bob")
+                        .EndTuple()
+                    .EndList().Build()
+                .Build();
+
+            const auto settings = TExecDataQuerySettings()
+                .CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), params, settings).GetValueSync();
+
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            CompareYson(R"([[[3500u];["None"];[1u];["Anna"]]])", FormatResultSetYson(result.GetResultSet(0)));
+
+            const Ydb::TableStats::QueryStats stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_EQUAL_C(2, stats.query_phases_size(), stats.DebugString());
+        }
+    }
 }
 
 } // namespace NKqp