[NEW RBO] Add support for multi consumer stages. (#31938)

5 files changed, 186 insertions, 39 deletions

diff --git a/ydb/core/kqp/opt/rbo/kqp_convert_to_physical.cpp b/ydb/core/kqp/opt/rbo/kqp_convert_to_physical.cpp
index 4cc679dab73..a8a160e91c0 100644
--- a/ydb/core/kqp/opt/rbo/kqp_convert_to_physical.cpp
+++ b/ydb/core/kqp/opt/rbo/kqp_convert_to_physical.cpp
@@ -48,6 +48,34 @@ TString GetValidJoinKind(const TString& joinKind) {
     return joinKind;
 }
 
+TExprNode::TPtr BuildMultiConsumerHandler(TExprNode::TPtr input, const ui32 numConsumers, TExprContext& ctx, TPositionHandle pos) {
+    TVector<TExprBase> branches;
+    auto inputIndex = NDq::BuildAtomList("0", pos, ctx);
+    for (ui32 i = 0; i < numConsumers; ++i) {
+        branches.emplace_back(inputIndex);
+        // Just an empty lambda.
+        // clang-format off
+        auto lambda = Build<TCoLambda>(ctx, pos)
+            .Args({"arg"})
+            .Body("arg")
+        .Done();
+        // clang-format on
+        branches.push_back(lambda);
+    }
+
+    // clang-format off
+    return Build<TCoSwitch>(ctx, pos)
+        .Input(input)
+        .BufferBytes()
+            .Value(ToString(128_MB))
+        .Build()
+        .FreeArgs()
+            .Add(branches)
+        .Build()
+     .Done().Ptr();
+     // clang-format on
+}
+
 TExprNode::TPtr ReplaceArg(TExprNode::TPtr input, TExprNode::TPtr arg, TExprContext &ctx, bool removeAliases = false) {
     // FIXME: This is not always correct, for example:
     // lambda($arg) { $val = expr($arg); return member($val `name)}
@@ -184,6 +212,10 @@ TExprNode::TPtr PeepholeStageLambda(TExprNode::TPtr stageLambda, TVector<TExprNo
     return TKqpProgram(newProgram).Lambda().Ptr();
 }
 
+bool IsMultiConsumerHandlerNeeded(const std::shared_ptr<IOperator>& op) {
+    return op->Props.NumOfConsumers.has_value() && op->Props.NumOfConsumers.value() > 1;
+}
+
 TExprNode::TPtr BuildCrossJoin(TOpJoin &join, TExprNode::TPtr leftInput, TExprNode::TPtr rightInput, TExprContext &ctx,
                                TPositionHandle pos) {
 
@@ -191,7 +223,7 @@ TExprNode::TPtr BuildCrossJoin(TOpJoin &join, TExprNode::TPtr leftInput, TExprNo
     TCoArgument rightArg{ctx.NewArgument(pos, "_kqp_right")};
 
     TVector<TExprNode::TPtr> keys;
-    for (auto iu : join.GetLeftInput()->GetOutputIUs()) {
+    for (const auto& iu : join.GetLeftInput()->GetOutputIUs()) {
         YQL_CLOG(TRACE, CoreDq) << "Converting Cross Join, left key: " << iu.GetFullName();
 
         // clang-format off
@@ -575,12 +607,12 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot &root, TRBOContext& rboCtx) {
 
                     // clang-format off
                     auto olapRead = Build<TKqpBlockReadOlapTableRanges>(ctx, op->Pos)
-                            .Table(opSource->TableCallable)
-                            .Ranges<TCoVoid>().Build()
-                            .Columns().Add(columns).Build()
-                            .Settings<TCoNameValueTupleList>().Build()
-                            .ExplainPrompt<TCoNameValueTupleList>().Build()
-                            .Process(processLambda)
+                        .Table(opSource->TableCallable)
+                        .Ranges<TCoVoid>().Build()
+                        .Columns().Add(columns).Build()
+                        .Settings<TCoNameValueTupleList>().Build()
+                        .ExplainPrompt<TCoNameValueTupleList>().Build()
+                        .Process(processLambda)
                     .Done().Ptr();
 
                     auto flowNonBlockRead = Build<TCoToFlow>(ctx, op->Pos)
@@ -599,6 +631,10 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot &root, TRBOContext& rboCtx) {
                         .Input(narrowMap)
                     .Done().Ptr();
                     // clang-format on
+
+                    if (IsMultiConsumerHandlerNeeded(op)) {
+                        currentStageBody = BuildMultiConsumerHandler(currentStageBody, op->Props.NumOfConsumers.value(), ctx, op->Pos);
+                    }
                     break;
                 }
                 default:
@@ -741,7 +777,11 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot &root, TRBOContext& rboCtx) {
                     .Build()
                 .Build()
             .Done().Ptr();
-            // clang-fort on
+            // clang-format on
+
+            if (IsMultiConsumerHandlerNeeded(op)) {
+               currentStageBody = BuildMultiConsumerHandler(currentStageBody, op->Props.NumOfConsumers.value(), ctx, op->Pos);
+            }
 
             stages[opStageId] = currentStageBody;
             stagePos[opStageId] = op->Pos;
@@ -884,17 +924,25 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot &root, TRBOContext& rboCtx) {
         YQL_CLOG(TRACE, CoreDq) << "Finalizing stage " << id;
 
         TVector<TExprNode::TPtr> inputs;
+        THashSet<int> processedInputsIds;
         for (const auto inputStageId : stageInputIds.at(id)) {
+            if (processedInputsIds.contains(inputStageId)) {
+                continue;
+            }
+            processedInputsIds.insert(inputStageId);
+
             auto inputStage = finalizedStages.at(inputStageId);
-            auto connection = graph.GetConnection(inputStageId, id);
-            YQL_CLOG(TRACE, CoreDq) << "Building connection: " << inputStageId << "->" << id << ", " << connection->Type;
-            TExprNode::TPtr newStage;
-            auto dqConnection = connection->BuildConnection(inputStage, stagePos.at(inputStageId), newStage, ctx);
-            if (newStage) {
-                txStages.push_back(newStage);
+            const auto connections = graph.GetConnections(inputStageId, id);
+            for (const auto& connection : connections) {
+                YQL_CLOG(TRACE, CoreDq) << "Building connection: " << inputStageId << "->" << id << ", " << connection->Type;
+                TExprNode::TPtr newStage;
+                auto dqConnection = connection->BuildConnection(inputStage, stagePos.at(inputStageId), newStage, ctx);
+                if (newStage) {
+                    txStages.push_back(newStage);
+                }
+                YQL_CLOG(TRACE, CoreDq) << "Built connection: " << inputStageId << "->" << id << ", " << connection->Type;
+                inputs.push_back(dqConnection);
             }
-            YQL_CLOG(TRACE, CoreDq) << "Built connection: " << inputStageId << "->" << id << ", " << connection->Type;
-            inputs.push_back(dqConnection);
         }
 
         TExprNode::TPtr stage;
diff --git a/ydb/core/kqp/opt/rbo/kqp_operator.h b/ydb/core/kqp/opt/rbo/kqp_operator.h
index e139a3403c2..070edc39a0d 100644
--- a/ydb/core/kqp/opt/rbo/kqp_operator.h
+++ b/ydb/core/kqp/opt/rbo/kqp_operator.h
@@ -122,6 +122,7 @@ struct TPhysicalOpProps {
     std::optional<int> StageId;
     std::optional<TString> Algorithm;
     std::optional<TOrderEnforcer> OrderEnforcer;
+    std::optional<ui32> NumOfConsumers;
     bool EnsureAtMostOne = false;
 
     std::optional<TRBOMetadata> Metadata;
@@ -140,13 +141,12 @@ struct TConnection {
         , FromSourceStageStorageType(fromSourceStageStorageType)
         , OutputIndex(outputIndex) {
     }
+    virtual ~TConnection() = default;
 
     virtual TExprNode::TPtr BuildConnection(TExprNode::TPtr inputStage, TPositionHandle pos, TExprNode::TPtr& newStage, TExprContext& ctx) = 0;
-
     template <typename T>
     TExprNode::TPtr BuildConnectionImpl(TExprNode::TPtr inputStage, TPositionHandle pos, TExprNode::TPtr& newStage, TExprContext& ctx);
-
-    virtual ~TConnection() = default;
+    ui32 GetOutputIndex() const { return OutputIndex; }
 
     TString Type;
     NYql::EStorageType FromSourceStageStorageType;
@@ -229,7 +229,8 @@ struct TStageGraph {
     THashMap<int, TSourceStageTraits> SourceStageRenames;
     THashMap<int, TVector<int>> StageInputs;
     THashMap<int, TVector<int>> StageOutputs;
-    THashMap<std::pair<int, int>, std::shared_ptr<TConnection>> Connections;
+    THashMap<std::pair<int, int>, TVector<std::shared_ptr<TConnection>>> Connections;
+    THashMap<int, int> StageOutputIndices;
 
     int AddStage() {
         int newStageId = StageIds.size();
@@ -273,15 +274,15 @@ struct TStageGraph {
         return NYql::EStorageType::NA;
     }
 
-    void Connect(int from, int to, std::shared_ptr<TConnection> conn) {
+    void Connect(int from, int to, std::shared_ptr<TConnection> connection) {
         auto &outputs = StageOutputs.at(from);
         outputs.push_back(to);
         auto &inputs = StageInputs.at(to);
         inputs.push_back(from);
-        Connections[std::make_pair(from, to)] = conn;
+        Connections[std::make_pair(from, to)].push_back(connection);
     }
 
-    std::shared_ptr<TConnection> GetConnection(int from, int to) { return Connections.at(std::make_pair(from, to)); }
+    TVector<std::shared_ptr<TConnection>> GetConnections(int from, int to) { return Connections.at(std::make_pair(from, to)); }
 
     /**
      * Generate an expression for stage inputs
@@ -290,6 +291,18 @@ struct TStageGraph {
     std::pair<TExprNode::TPtr, TExprNode::TPtr> GenerateStageInput(int &stageInputCounter, TExprNode::TPtr &node, TExprContext &ctx,
                                                                    int fromStage);
 
+    ui32 GetOutputIndex(ui32 stageIndex) {
+        ui32 outputIndex{0};
+        auto it = StageOutputIndices.find(stageIndex);
+        if (it != StageOutputIndices.end()) {
+            it->second++;
+            outputIndex = it->second;
+        } else {
+            StageOutputIndices[stageIndex] = 0;
+        }
+        return outputIndex;
+    }
+
     void TopologicalSort();
 private:
 
diff --git a/ydb/core/kqp/opt/rbo/kqp_rbo_rules.cpp b/ydb/core/kqp/opt/rbo/kqp_rbo_rules.cpp
index f6333e102f0..db3be0891a1 100644
--- a/ydb/core/kqp/opt/rbo/kqp_rbo_rules.cpp
+++ b/ydb/core/kqp/opt/rbo/kqp_rbo_rules.cpp
@@ -110,7 +110,7 @@ TExprNode::TPtr BuildFilterLambdaFromConjuncts(TPositionHandle pos, TVector<TFil
     }
 }
 
-THashSet<TString> CmpOperators{"=", "<", ">", "<=", ">="};
+const THashSet<TString> CmpOperators{"=", "<", ">", "<=", ">="};
 
 TExprNode::TPtr PruneCast(TExprNode::TPtr node) {
     if (node->IsCallable("ToPg") || node->IsCallable("PgCast")) {
@@ -137,6 +137,15 @@ TVector<TInfoUnit> GetHashableKeys(const std::shared_ptr<IOperator> &input) {
     return hashableKeys;
 }
 
+void UpdateNumOfConsumers(std::shared_ptr<IOperator> &input) {
+    auto &props = input->Props;
+    if (props.NumOfConsumers.has_value()) {
+        props.NumOfConsumers.value() += 1;
+    } else {
+        props.NumOfConsumers = 1;
+    }
+}
+
 bool IsNullRejectingPredicate(const TFilterInfo &filter, TExprContext &ctx) {
     Y_UNUSED(ctx);
 #ifdef DEBUG_PREDICATE
@@ -1173,6 +1182,7 @@ bool TAssignStagesRule::MatchAndApply(std::shared_ptr<IOperator> &input, TRBOCon
     } else if (input->Kind == EOperator::Filter || input->Kind == EOperator::Map) {
         auto childOp = CastOperator<IUnaryOperator>(input)->GetInput();
         auto prevStageId = *(childOp->Props.StageId);
+        UpdateNumOfConsumers(childOp);
 
         // If the child operator is a source, it requires its own stage
         // So we have build a new stage for current operator
@@ -1188,7 +1198,7 @@ bool TAssignStagesRule::MatchAndApply(std::shared_ptr<IOperator> &input, TRBOCon
                     break;
                 }
                 case NYql::EStorageType::ColumnStorage: {
-                    connection.reset(new TUnionAllConnection(NYql::EStorageType::ColumnStorage));
+                    connection.reset(new TUnionAllConnection(NYql::EStorageType::ColumnStorage, props.StageGraph.GetOutputIndex(prevStageId)));
                     break;
                 }
                 default: {
@@ -1225,6 +1235,8 @@ bool TAssignStagesRule::MatchAndApply(std::shared_ptr<IOperator> &input, TRBOCon
         props.StageGraph.Connect(prevStageId, newStageId,conn);
     } else if (input->Kind == EOperator::UnionAll) {
         auto unionAll = CastOperator<TOpUnionAll>(input);
+        UpdateNumOfConsumers(unionAll->GetLeftInput());
+        UpdateNumOfConsumers(unionAll->GetRightInput());
 
         auto leftStage = unionAll->GetLeftInput()->Props.StageId;
         auto rightStage = unionAll->GetRightInput()->Props.StageId;
@@ -1232,10 +1244,12 @@ bool TAssignStagesRule::MatchAndApply(std::shared_ptr<IOperator> &input, TRBOCon
         auto newStageId = props.StageGraph.AddStage();
         unionAll->Props.StageId = newStageId;
 
-        props.StageGraph.Connect(*leftStage, newStageId,
-                                 std::make_shared<TUnionAllConnection>(props.StageGraph.GetStorageType(*leftStage)));
-        props.StageGraph.Connect(*rightStage, newStageId,
-                                 std::make_shared<TUnionAllConnection>(props.StageGraph.GetStorageType(*rightStage)));
+        props.StageGraph.Connect(
+            *leftStage, newStageId,
+            std::make_shared<TUnionAllConnection>(props.StageGraph.GetStorageType(*leftStage), props.StageGraph.GetOutputIndex(*leftStage)));
+        props.StageGraph.Connect(
+            *rightStage, newStageId,
+            std::make_shared<TUnionAllConnection>(props.StageGraph.GetStorageType(*rightStage), props.StageGraph.GetOutputIndex(*rightStage)));
 
         YQL_CLOG(TRACE, CoreDq) << "Assign stages union_all";
     } else if (input->Kind == EOperator::Aggregate) {
diff --git a/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp b/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp
index dc5dff5eb87..817ac097ce1 100644
--- a/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp
+++ b/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp
@@ -725,7 +725,7 @@ TExprNode::TPtr RewriteSelect(const TExprNode::TPtr &node, TExprContext &ctx, co
                         .Alias(alias)
                         .Columns(readExpr.Columns())
                         .SourceType(GetTableSourceType(tableDesc, ctx, node->Pos()))
-                        .UniqueId().Value(std::to_string(uniqueSourceIdCounter++)).Build()
+                        .UniqueId().Value(std::to_string(uniqueSourceIdCounter)).Build()
                     .Done().Ptr();
                     // clang-format on
                 }
diff --git a/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp b/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp
index 2ca72ade58e..4cb64ed59bc 100644
--- a/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp
+++ b/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp
@@ -576,6 +576,23 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         UNIT_ASSERT_C(resultUpsert.IsSuccess(), resultUpsert.GetIssues().ToString());
     }
 
+    void InsertIntoSchema1(NYdb::NTable::TTableClient &db, std::string tableName, int numRows) {
+        NYdb::TValueBuilder rows;
+        rows.BeginList();
+        for (size_t i = 0; i < numRows; ++i) {
+            rows.AddListItem()
+                .BeginStruct()
+                .AddMember("a").Int64(i)
+                .AddMember("b").Int64(i + 2)
+                .AddMember("c").Int64(i + 1)
+                .EndStruct();
+        }
+        rows.EndList();
+        auto resultUpsert = db.BulkUpsert(tableName, rows.Build()).GetValueSync();
+        UNIT_ASSERT_C(resultUpsert.IsSuccess(), resultUpsert.GetIssues().ToString());
+    }
+
+ 
     void CreateSimpleTable(TKikimrRunner &kikimr) {
         auto db = kikimr.GetTableClient();
         auto session = db.CreateSession().GetValueSync().GetSession();
@@ -742,14 +759,6 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
             R"(
                 --!syntax_pg
                 SET TablePathPrefix = "/Root/";
-                select sum(t1.c) as sum from t1 group by t1.b
-                union all
-                select sum(t1.b) as sum from t1
-                order by sum;
-            )",
-            R"(
-                --!syntax_pg
-                SET TablePathPrefix = "/Root/";
                 select t1.b, min(t1.a) from t1 group by t1.b order by t1.b;
             )",
             R"(
@@ -900,7 +909,6 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         std::vector<std::string> results = {
             R"([["1";"4"];["2";"6"]])",
             R"([["1";"4"];["2";"6"]])",
-            R"([["4"];["6"];["8"]])",
             R"([["1";"1"];["2";"0"]])",
             R"([["1";"3"];["2";"4"]])",
             R"([["1";"2"];["2";"3"]])",
@@ -943,6 +951,70 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         TestAggregation(ColumnStore);
     }
 
+    Y_UNIT_TEST(MultiConsumer) {
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetEnableNewRBO(true);
+        appConfig.MutableTableServiceConfig()->SetAllowOlapDataQuery(true);
+        appConfig.MutableTableServiceConfig()->SetEnableFallbackToYqlOptimizer(false);
+        TKikimrRunner kikimr(NKqp::TKikimrSettings(appConfig).SetWithSampleTables(false));
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        session.ExecuteSchemeQuery(R"(
+            CREATE TABLE `/Root/t1` (
+                a Int64 NOT NULL,
+	            b Int64,
+                c Int64,
+                primary key(a)
+            ) with (Store = Column);
+
+            CREATE TABLE `/Root/t2` (
+                a Int64 NOT NULL,
+	            b Int64,
+                c Int64,
+                primary key(a)
+            ) with (Store = Column);
+        )").GetValueSync();
+
+        db = kikimr.GetTableClient();
+        auto session2 = db.CreateSession().GetValueSync().GetSession();
+        const std::vector<std::pair<std::string, int>> tables{{"/Root/t1", 4}, {"/Root/t2", 3}};
+        for (const auto &[table, rowsNum] : tables) {
+            InsertIntoSchema1(db, table, rowsNum);
+        }
+
+        std::vector<std::string> queries = {
+            R"(
+                --!syntax_pg
+                set TablePathPrefix = "/Root/";
+                select t1.a as a from t1
+                union all
+                select t1.a as a from t1 order by a;
+            )",
+            R"(
+                --!syntax_pg
+                SET TablePathPrefix = "/Root/";
+                select sum(t1.c) as sum from t1 group by t1.b
+                union all
+                select sum(t1.c) as sum from t1
+                order by sum;
+            )"
+        };
+
+        std::vector<std::string> results = {
+            R"([["0"];["0"];["1"];["1"];["2"];["2"];["3"];["3"]])",
+            R"([["1"];["2"];["3"];["4"];["10"]])"
+        };
+
+        for (ui32 i = 0; i < queries.size(); ++i) {
+            const auto &query = queries[i];
+            auto result = session2.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            UNIT_ASSERT_VALUES_EQUAL(FormatResultSetYson(result.GetResultSet(0)), results[i]);
+            //Cout << FormatResultSetYson(result.GetResultSet(0)) << Endl;
+        }
+    }
+
     Y_UNIT_TEST(UnionAll) {
         NKikimrConfig::TAppConfig appConfig;
         appConfig.MutableTableServiceConfig()->SetEnableNewRBO(true);