[KQP] Add new rule to KQP pipeline (#20809)

denis0x0D · web-flow · commit bfb3eb046c52 · 2025-07-09T09:50:43.000+03:00
diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
@@ -103,6 +103,7 @@ class TKqpPhysicalOptTransformer : public TOptimizeTransformerBase {
         AddHandler(1, &TCoExtractMembers::Match, HNDL(PushExtractMembersToStage<true>));
         AddHandler(1, &TCoFlatMapBase::Match, HNDL(BuildFlatmapStage<true>));
         AddHandler(1, &TCoCombineByKey::Match, HNDL(PushCombineToStage<true>));
+        AddHandler(1, &TCoCombineByKey::Match, HNDL(PushCombineToStageDependsOnOtherStage<true>));
         AddHandler(1, &TCoPartitionsByKeys::Match, HNDL(BuildPartitionsStage<true>));
         AddHandler(1, &TCoFinalizeByKey::Match, HNDL(BuildFinalizeByKeyStage<true>));
         AddHandler(1, &TCoShuffleByKeys::Match, HNDL(BuildShuffleStage<true>));
@@ -330,6 +331,15 @@ class TKqpPhysicalOptTransformer : public TOptimizeTransformerBase {
         return output;
     }
 
+    template <bool IsGlobal>
+    TMaybeNode<TExprBase> PushCombineToStageDependsOnOtherStage(TExprBase node, TExprContext& ctx,
+        IOptimizationContext& optCtx, const TGetParents& getParents)
+    {
+        TExprBase output = DqPushCombineToStageDependsOnOtherStage(node, ctx, optCtx, *getParents(), IsGlobal);
+        DumpAppliedRule("PushCombineToStageDependsOnOtherStage", node.Ptr(), output.Ptr(), ctx);
+        return output;
+    }
+
     template <bool IsGlobal>
     TMaybeNode<TExprBase> BuildShuffleStage(TExprBase node, TExprContext& ctx,
         IOptimizationContext& optCtx, const TGetParents& getParents)
diff --git a/ydb/core/kqp/ut/opt/kqp_agg_ut.cpp b/ydb/core/kqp/ut/opt/kqp_agg_ut.cpp
@@ -244,6 +244,88 @@ Y_UNIT_TEST_SUITE(KqpAgg) {
         }
        
     }
+
+    Y_UNIT_TEST(AggWithSqlIn) {
+        auto settings = TKikimrSettings().SetWithSampleTables(false);
+        TKikimrRunner kikimr(settings);
+
+        auto tableClient = kikimr.GetTableClient();
+        auto session = tableClient.CreateSession().GetValueSync().GetSession();
+
+        auto queryClient = kikimr.GetQueryClient();
+        auto result = queryClient.GetSession().GetValueSync();
+        NStatusHelpers::ThrowOnError(result);
+        auto session2 = result.GetSession();
+
+        auto res = session.ExecuteSchemeQuery(R"(
+            CREATE TABLE `/Root/t1` (
+                a Int64	NOT NULL,
+                b Int32,
+                primary key(a)
+            )
+            PARTITION BY HASH(a)
+            WITH (STORE = COLUMN);
+        )").GetValueSync();
+        UNIT_ASSERT(res.IsSuccess());
+
+       res = session.ExecuteSchemeQuery(R"(
+            CREATE TABLE `/Root/t2` (
+                a Int64	NOT NULL,
+                b Int32,
+                primary key(a)
+            )
+            PARTITION BY HASH(a)
+            WITH (STORE = COLUMN);
+        )").GetValueSync();
+        UNIT_ASSERT(res.IsSuccess());
+
+        res = session.ExecuteSchemeQuery(R"(
+            CREATE TABLE `/Root/t3` (
+                a Int64	NOT NULL,
+                b Int32,
+                primary key(a)
+            )
+            PARTITION BY HASH(a)
+            WITH (STORE = COLUMN);
+        )").GetValueSync();
+        UNIT_ASSERT(res.IsSuccess());
+
+        auto insertRes = session2.ExecuteQuery(R"(
+            INSERT INTO `/Root/t1` (a, b) VALUES (1, 1);
+            INSERT INTO `/Root/t2` (a, b) VALUES (1, 1);
+            INSERT INTO `/Root/t3` (a, b) VALUES (1, 1);
+            INSERT INTO `/Root/t1` (a, b) VALUES (2, 1);
+            INSERT INTO `/Root/t2` (a, b) VALUES (2, 1);
+            INSERT INTO `/Root/t3` (a, b) VALUES (2, 1);
+            INSERT INTO `/Root/t1` (a, b) VALUES (3, 1);
+            INSERT INTO `/Root/t2` (a, b) VALUES (3, 1);
+            INSERT INTO `/Root/t3` (a, b) VALUES (3, 1);
+        )", NYdb::NQuery::TTxControl::NoTx()).GetValueSync();
+        UNIT_ASSERT(insertRes.IsSuccess());
+
+        std::vector<TString> queries = {
+            R"(
+                SELECT sum(CASE WHEN (t1.a IN (SELECT t2.a FROM t2)) THEN 1 ELSE 0 END),
+                       sum(CASE WHEN (t1.a IN (SELECT t3.a FROM t3)) THEN 1 ELSE 0 END)
+                FROM t1;
+            )",
+        };
+
+        for (ui32 i = 0; i < queries.size(); ++i) {
+            const auto query = queries[i];
+            auto result =
+                session2
+                    .ExecuteQuery(query, NYdb::NQuery::TTxControl::NoTx(), NYdb::NQuery::TExecuteQuerySettings().ExecMode(NQuery::EExecMode::Explain))
+                    .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+
+            result = session2.ExecuteQuery(query, NYdb::NQuery::TTxControl::NoTx(), NYdb::NQuery::TExecuteQuerySettings()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+
+            TString output = FormatResultSetYson(result.GetResultSet(0));
+            UNIT_ASSERT_VALUES_EQUAL(FormatResultSetYson(result.GetResultSet(0)), "[[[3];[3]]]");
+        }
+    }
 }
 
 } // namespace NKikimr::NKqp
diff --git a/ydb/library/yql/dq/opt/dq_opt.cpp b/ydb/library/yql/dq/opt/dq_opt.cpp
@@ -140,6 +140,15 @@ bool IsDqDependsOnStage(const TExprBase& node, const TDqStageBase& stage) {
     });
 }
 
+bool IsDqDependsOnOtherStage(const TExprBase& node, const TDqStageBase& stage) {
+    return !!FindNode(node.Ptr(), [ptr = stage.Raw()](const TExprNode::TPtr& node) {
+        if (TMaybeNode<TDqStage>(node)) {
+            return node.Get() != ptr;
+        }
+        return false;
+    });
+}
+
 bool IsDqDependsOnStageOutput(const TExprBase& node, const TDqStageBase& stage, ui32 outputIndex) {
     return !!FindNode(node.Ptr(), [ptr = stage.Raw(), outputIndex](const TExprNode::TPtr& exprNode) {
         if (TDqOutput::Match(exprNode.Get())) {
diff --git a/ydb/library/yql/dq/opt/dq_opt.h b/ydb/library/yql/dq/opt/dq_opt.h
@@ -38,6 +38,7 @@ inline bool IsDqCompletePureExpr(const NNodes::TExprBase& node, bool isPrecomput
 
 bool IsDqSelfContainedExpr(const NNodes::TExprBase& node);
 bool IsDqDependsOnStage(const NNodes::TExprBase& node, const NNodes::TDqStageBase& stage);
+bool IsDqDependsOnOtherStage(const NNodes::TExprBase& node, const NNodes::TDqStageBase& stage);
 bool IsDqDependsOnStageOutput(const NNodes::TExprBase& node, const NNodes::TDqStageBase& stage, ui32 outputIndex);
 
 bool CanPushDqExpr(const NNodes::TExprBase& expr, const NNodes::TDqStageBase& stage);
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -1075,6 +1075,8 @@ TExprBase DqBuildLMapOverMuxStage(TExprBase node, TExprContext& ctx, IOptimizati
     return DqBuildLMapOverMuxStageStub<TCoLMap>(node, ctx, optCtx, parentsMap);
 }
 
+
+
 TExprBase DqPushCombineToStage(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage)
 {
@@ -1156,6 +1158,112 @@ TExprBase DqPushCombineToStage(TExprBase node, TExprContext& ctx, IOptimizationC
     return result.Cast();
 }
 
+TExprBase DqPushCombineToStageDependsOnOtherStage(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx, const TParentsMap& parentsMap,
+                                                  bool allowStageMultiUsage) {
+    Y_UNUSED(optCtx);
+    if (!node.Maybe<TCoCombineByKey>().Input().Maybe<TDqCnUnionAll>()) {
+        return node;
+    }
+
+    auto combine = node.Cast<TCoCombineByKey>();
+    auto dqUnion = combine.Input().Cast<TDqCnUnionAll>();
+
+    if (!IsSingleConsumerConnection(dqUnion, parentsMap, allowStageMultiUsage)) {
+        return node;
+    }
+
+    if (IsDqDependsOnOtherStage(combine.InitHandlerLambda(), dqUnion.Output().Stage()) &&
+        IsDqDependsOnOtherStage(combine.UpdateHandlerLambda(), dqUnion.Output().Stage()) &&
+        !IsDqDependsOnOtherStage(combine.PreMapLambda(), dqUnion.Output().Stage()) &&
+        !IsDqDependsOnOtherStage(combine.KeySelectorLambda(), dqUnion.Output().Stage()) &&
+        !IsDqDependsOnOtherStage(combine.FinishHandlerLambda(), dqUnion.Output().Stage())) {
+
+        // Collect all connections for `UpdateHnadler` and `InitHandler` lambda, we want to replace them.
+        TExprNode::TListType connections{dqUnion.Ptr()};
+        auto connectionPredicate = [](const TExprNode::TPtr& node) { return !!TMaybeNode<TDqConnection>(node); };
+        const auto connectionsInitHandler = FindNodes(combine.InitHandlerLambda().Ptr(), connectionPredicate);
+        const auto connectionsUpdateHandler = FindNodes(combine.UpdateHandlerLambda().Ptr(), connectionPredicate);
+
+        if (connectionsInitHandler.empty() || (connectionsInitHandler.size() != connectionsUpdateHandler.size())) {
+            return node;
+        }
+
+        // Check that all collected connections are the same.
+        for (ui32 i = 0; i < connectionsInitHandler.size(); ++i) {
+            if ((connectionsInitHandler[i].Get() != connectionsUpdateHandler[i].Get()) || (connectionsInitHandler[i].Get() == dqUnion.Raw()) ||
+                (connectionsUpdateHandler[i].Get() == dqUnion.Raw())) {
+                return node;
+            }
+        }
+
+        connections.insert(connections.end(), connectionsInitHandler.begin(), connectionsInitHandler.end());
+
+        // Arguments for DqStage.
+        TVector<TCoArgument> inputArgs;
+        for (ui32 i = 0; i < connections.size(); ++i) {
+            TCoArgument arg = Build<TCoArgument>(ctx, node.Pos())
+                .Name(TStringBuilder() << "input_arg_" << i)
+                .Done();
+            inputArgs.push_back(arg);
+        }
+
+        // Arguments for the Flatmap's lambdas.
+        TVector<TCoArgument> lambdaArgs;
+        for (ui32 i = 1; i < connections.size(); ++i) {
+            TCoArgument lambdaArg = Build<TCoArgument>(ctx, node.Pos())
+                .Name(TStringBuilder() << "lambda_arg_" << i)
+                .Done();
+            lambdaArgs.push_back(lambdaArg);
+        }
+
+        TNodeOnNodeOwnedMap replaces;
+        replaces[connections[0].Get()] = inputArgs.front().Ptr();
+        for (ui32 i = 0; i < lambdaArgs.size(); ++i) {
+            replaces[connections[i + 1].Get()] = lambdaArgs[i].Ptr();
+        }
+
+        // Replace connections with `lambdaArgs` for the future flatmaps.
+        auto newBody = TExprBase(ctx.ReplaceNodes(combine.Ptr(), replaces));
+
+        // For input args in range [1, n] wrap to `SqueezeToList` and create
+        // a flatmap.
+        for (ui32 i = 1; i < inputArgs.size(); ++i) {
+            auto squeezeToListArg = Build<TCoSqueezeToList>(ctx, node.Pos())
+                .Stream(inputArgs[i])
+            .Done();
+
+            newBody = Build<TCoFlatMap>(ctx, node.Pos())
+                .Input(squeezeToListArg)
+                .Lambda()
+                    .Args(lambdaArgs[i - 1])
+                    .Body(newBody)
+                    .Build()
+                .Done();
+        }
+
+        // Build a stage with inputs.
+        auto stage = Build<TDqStage>(ctx, node.Pos())
+            .Inputs()
+                .Add(connections)
+                .Build()
+            .Program()
+                .Args(inputArgs)
+                .Body(newBody)
+                .Build()
+            .Settings(TDqStageSettings().BuildNode(ctx, node.Pos()))
+            .Done();
+
+        return Build<TDqCnUnionAll>(ctx, node.Pos())
+            .Output()
+                .Stage(stage)
+                .Index().Build("0")
+                .Build()
+            .Done();
+    }
+
+    return node;
+}
+
 NNodes::TExprBase DqPushAggregateCombineToStage(NNodes::TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage)
 {
@@ -1169,6 +1277,10 @@ NNodes::TExprBase DqPushAggregateCombineToStage(NNodes::TExprBase node, TExprCon
         return node;
     }
 
+    if (!IsDqCompletePureExpr(aggCombine.Handlers())) {
+        return node;
+    }
+
     auto lambda = Build<TCoLambda>(ctx, aggCombine.Pos())
             .Args({"stream"})
             .Body<TCoAggregateCombine>()
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.h b/ydb/library/yql/dq/opt/dq_opt_phy.h
@@ -60,6 +60,9 @@ NNodes::TExprBase DqPushFlatmapToStage(NNodes::TExprBase node, TExprContext& ctx
 NNodes::TExprBase DqPushCombineToStage(NNodes::TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage = true);
 
+NNodes::TExprBase DqPushCombineToStageDependsOnOtherStage(NNodes::TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
+    const TParentsMap& parentsMap, bool allowStageMultiUsage = true);
+
 NNodes::TExprBase DqPushAggregateCombineToStage(NNodes::TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage = true);
 
