[NEW RBO] Show BlockHashJoin in explain (#46362)

9 files changed, 81 insertions, 58 deletions

diff --git a/ydb/core/kqp/opt/rbo/kqp_operator.cpp b/ydb/core/kqp/opt/rbo/kqp_operator.cpp
index 32ac29917d9..82036e02302 100644
--- a/ydb/core/kqp/opt/rbo/kqp_operator.cpp
+++ b/ydb/core/kqp/opt/rbo/kqp_operator.cpp
@@ -750,8 +750,16 @@ TString GetJoinAlgoName(NKqp::EJoinAlgoType joinAlgo) {
     return "Unknown";
 }
 
-TString GetExplainJoinAlgoName(NKqp::EJoinAlgoType joinAlgo) {
-    if (joinAlgo == NKqp::EJoinAlgoType::GraceJoin) {
+TString GetExplainJoinAlgoName(const TPhysicalOpProps& props) {
+    Y_ENSURE(props.JoinAlgo.has_value(), "Join algorithm has not been selected");
+    Y_ENSURE(props.UseBlockHashJoin.has_value(), "Physical join implementation has not been selected");
+
+    if (*props.UseBlockHashJoin) {
+        return "BlockHash";
+    }
+
+    const auto joinAlgo = *props.JoinAlgo;
+    if (joinAlgo == NKqp::EJoinAlgoType::GraceJoin || joinAlgo == NKqp::EJoinAlgoType::ReverseBlockJoin) {
         return "Grace";
     }
     return GetJoinAlgoName(joinAlgo);
@@ -798,8 +806,7 @@ static TString FormatJoinKeys(const TVector<std::pair<TInfoUnit, TInfoUnit>>& jo
 
 NJson::TJsonValue TOpJoin::ToJson(ui32 explainFlags) {
     auto res = IOperator::ToJson(explainFlags);
-    const auto joinAlgo = Props.JoinAlgo.value_or(NKqp::EJoinAlgoType::Undefined);
-    const auto joinAlgoName = GetExplainJoinAlgoName(joinAlgo);
+    const auto joinAlgoName = GetExplainJoinAlgoName(Props);
 
     if (JoinKind == "Cross") {
         res["Name"] = "CrossJoin";
diff --git a/ydb/core/kqp/opt/rbo/kqp_operator.h b/ydb/core/kqp/opt/rbo/kqp_operator.h
index a0e3c0dd928..88d6a5898a8 100644
--- a/ydb/core/kqp/opt/rbo/kqp_operator.h
+++ b/ydb/core/kqp/opt/rbo/kqp_operator.h
@@ -125,6 +125,9 @@ struct TPhysicalOpProps {
     std::optional<TRBOMetadata> Metadata;
     std::optional<TRBOStatistics> Statistics;
     std::optional<NKikimr::NKqp::EJoinAlgoType> JoinAlgo;
+    // Resolved physical implementation for a join. std::nullopt means that
+    // physical join selection has not run yet.
+    std::optional<bool> UseBlockHashJoin;
     std::optional<double> Cost;
 
     // CBO decision for this join's input edges.
@@ -148,6 +151,7 @@ private:
         Metadata = other.Metadata;
         Statistics = other.Statistics;
         JoinAlgo = other.JoinAlgo;
+        UseBlockHashJoin = other.UseBlockHashJoin;
         Cost = other.Cost;
         LeftShuffleBy = other.LeftShuffleBy;
         RightShuffleBy = other.RightShuffleBy;
diff --git a/ydb/core/kqp/opt/rbo/kqp_rbo_utils.cpp b/ydb/core/kqp/opt/rbo/kqp_rbo_utils.cpp
index a5c4ca4a8a0..08ea4e141aa 100644
--- a/ydb/core/kqp/opt/rbo/kqp_rbo_utils.cpp
+++ b/ydb/core/kqp/opt/rbo/kqp_rbo_utils.cpp
@@ -90,6 +90,18 @@ bool JoinOutputsRight(const TString& joinKind) {
     return joinKind != "LeftOnly" && joinKind != "LeftSemi";
 }
 
+TString GetValidJoinKind(const TString& joinKind) {
+    const auto joinKindLowered = to_lower(joinKind);
+    if (joinKindLowered == "left") {
+        return "Left";
+    } else if (joinKindLowered == "inner") {
+        return "Inner";
+    } else if (joinKindLowered == "cross") {
+        return "Cross";
+    }
+    return joinKind;
+}
+
 TVector<TInfoUnit> IUSetDiff(TVector<TInfoUnit> left, TVector<TInfoUnit> right) {
     TVector<TInfoUnit> res;
     for (const auto& unit : left) {
diff --git a/ydb/core/kqp/opt/rbo/kqp_rbo_utils.h b/ydb/core/kqp/opt/rbo/kqp_rbo_utils.h
index 3c1e25028c4..55468ef0373 100644
--- a/ydb/core/kqp/opt/rbo/kqp_rbo_utils.h
+++ b/ydb/core/kqp/opt/rbo/kqp_rbo_utils.h
@@ -26,6 +26,7 @@ TVector<TInfoUnit> GetSubplanResultIUs(const TIntrusivePtr<IOperator>& op);
 
 bool JoinOutputsLeft(const TString& joinKind);
 bool JoinOutputsRight(const TString& joinKind);
+TString GetValidJoinKind(const TString& joinKind);
 
 TVector<TInfoUnit> IUSetDiff(TVector<TInfoUnit> left, TVector<TInfoUnit> right);
 TVector<TInfoUnit> IUSetIntersect(TVector<TInfoUnit> left, TVector<TInfoUnit> right);
diff --git a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_convert_to_physical.cpp b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_convert_to_physical.cpp
index a55f017c70c..9b363d088c9 100644
--- a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_convert_to_physical.cpp
+++ b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_convert_to_physical.cpp
@@ -164,13 +164,14 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot& root, TRBOContext& rboCtx) {
             YQL_CLOG(TRACE, CoreDq) << "Converted Sort " << opStageId;
         } else if (op->Kind == EOperator::Join) {
             auto join = CastOperator<TOpJoin>(op);
+            Y_ENSURE(join->Props.UseBlockHashJoin.has_value(), "Physical join implementation has not been selected");
 
             auto [leftArg, leftInput] = graph.GenerateStageInput(stageInputCounter, op->Pos, ctx);
             stageArgs[opStageId].push_back(leftArg);
             auto [rightArg, rightInput] = graph.GenerateStageInput(stageInputCounter, op->Pos, ctx);
             stageArgs[opStageId].push_back(rightArg);
 
-            currentStageBody = Build<TPhysicalJoinBuilder>(join, ctx, op->Pos, leftInput, rightInput, rboCtx.KqpCtx.Config->GetUseBlockHashJoin());
+            currentStageBody = Build<TPhysicalJoinBuilder>(join, ctx, op->Pos, leftInput, rightInput, *join->Props.UseBlockHashJoin);
 
             if (!join->IsSingleConsumer()) {
                 currentStageBody = NPhysicalConvertionUtils::BuildMultiConsumerHandler(currentStageBody, join->GetNumOfConsumers(), ctx, op->Pos);
diff --git a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.cpp b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.cpp
index 8570766ce49..1ebc1b7fe52 100644
--- a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.cpp
+++ b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.cpp
@@ -1,5 +1,6 @@
 #include "kqp_rbo_physical_join_builder.h"
 #include "kqp_rbo_physical_convertion_utils.h"
+#include <ydb/core/kqp/opt/rbo/kqp_rbo_utils.h>
 
 #include <yql/essentials/core/yql_expr_type_annotation.h>
 
@@ -8,18 +9,6 @@ using namespace NKikimr;
 using namespace NKikimr::NKqp;
 
 
-TString TPhysicalJoinBuilder::GetValidJoinKind(const TString& joinKind) const {
-    const auto joinKindLowered = to_lower(joinKind);
-    if (joinKindLowered == "left") {
-        return "Left";
-    } else if (joinKindLowered == "inner") {
-        return "Inner";
-    } else if (joinKindLowered == "cross") {
-        return "Cross";
-    }
-    return joinKind;
-}
-
 TExprNode::TPtr TPhysicalJoinBuilder::BuildCrossJoin(TExprNode::TPtr leftInput, TExprNode::TPtr rightInput) {
     TCoArgument leftArg{Ctx.NewArgument(Pos, "_kqp_left")};
     TCoArgument rightArg{Ctx.NewArgument(Pos, "_kqp_right")};
@@ -428,10 +417,6 @@ TExprNode::TPtr TPhysicalJoinBuilder::BuildPhysicalJoin(TExprNode::TPtr leftInpu
     Y_ENSURE(props.JoinAlgo.has_value());
     const auto joinAlgo = *(props.JoinAlgo);
 
-    useBlockHashJoin = useBlockHashJoin &&
-                       (joinAlgo == NKikimr::NKqp::EJoinAlgoType::GraceJoin || joinAlgo == NKikimr::NKqp::EJoinAlgoType::ReverseBlockJoin) &&
-                       (joinType == "Inner"sv || joinType == "Left"sv || joinType == "LeftSemi"sv || joinType == "LeftOnly"sv);
-
     const auto leftInputType = Join->GetLeftInput()->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
     const auto rightInputType = Join->GetRightInput()->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
     TModifyKeysList remapLeft;
diff --git a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.h b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.h
index b2708c7860b..1ba31f3d215 100644
--- a/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.h
+++ b/ydb/core/kqp/opt/rbo/physical_conversion/kqp_rbo_physical_join_builder.h
@@ -23,7 +23,6 @@ public:
 private:
     TExprNode::TPtr BuildPhysicalJoin(TExprNode::TPtr leftInput, TExprNode::TPtr rightInput, bool useBlockHashJoins);
     TExprNode::TPtr BuildCrossJoin(TExprNode::TPtr leftInput, TExprNode::TPtr rightInput);
-    TString GetValidJoinKind(const TString& joinKind) const;
     void PrepareJoinKeys(TVector<TString>& leftJoinKeys, TVector<TString>& rightJoinKeys, TModifyKeysList& remapLeft, TModifyKeysList& remapRight,
                          THashMap<TString, TString>& leftColumnRemap, THashMap<TString, TString>& rightColumnRemap, TVector<TString>& leftJoinKeyRenames,
                          TVector<TString>& rightJoinKeyRenames, const TStructExprType* leftInputType, const TStructExprType* rightInputType, const bool outer,
diff --git a/ydb/core/kqp/opt/rbo/rules/assign_stages.cpp b/ydb/core/kqp/opt/rbo/rules/assign_stages.cpp
index 2046cee02bb..b9e23e9e37a 100644
--- a/ydb/core/kqp/opt/rbo/rules/assign_stages.cpp
+++ b/ydb/core/kqp/opt/rbo/rules/assign_stages.cpp
@@ -1,4 +1,5 @@
 #include <ydb/core/kqp/opt/rbo/kqp_rbo_rules.h>
+#include <ydb/core/kqp/opt/rbo/kqp_rbo_utils.h>
 #include <ydb/core/kqp/provider/yql_kikimr_settings.h>
 
 namespace NKikimr::NKqp {
@@ -8,24 +9,27 @@ namespace {
 using namespace NKikimr;
 using namespace NKikimr::NKqp;
 
-void MaybeSetJoinAlgo(TPhysicalOpProps& props, const TRBOContext& rboCtx) {
-    if (props.JoinAlgo.has_value()) {
-        return;
-    }
-
-    auto joinMode = rboCtx.KqpCtx.Config->GetHashJoinMode();
-    NKikimr::NKqp::EJoinAlgoType joinAlgo;
-    switch (joinMode) {
-        case NYql::NDq::EHashJoinMode::Map: {
-            joinAlgo = NKikimr::NKqp::EJoinAlgoType::MapJoin;
-            break;
-        }
-        default: {
-            joinAlgo = NKikimr::NKqp::EJoinAlgoType::GraceJoin;
-            break;
+void FinalizeJoinPhysicalProps(TOpJoin& join, const TRBOContext& rboCtx) {
+    auto& props = join.Props;
+    if (!props.JoinAlgo.has_value()) {
+        const auto joinMode = rboCtx.KqpCtx.Config->GetHashJoinMode();
+        switch (joinMode) {
+            case NYql::NDq::EHashJoinMode::Map: {
+                props.JoinAlgo = EJoinAlgoType::MapJoin;
+                break;
+            }
+            default: {
+                props.JoinAlgo = EJoinAlgoType::GraceJoin;
+                break;
+            }
         }
     }
-    props.JoinAlgo = joinAlgo;
+
+    const auto joinKind = GetValidJoinKind(join.JoinKind);
+    const auto joinAlgo = *props.JoinAlgo;
+    props.UseBlockHashJoin = rboCtx.KqpCtx.Config->GetUseBlockHashJoin()
+        && (joinAlgo == EJoinAlgoType::GraceJoin || joinAlgo == EJoinAlgoType::ReverseBlockJoin)
+        && (joinKind == "Inner" || joinKind == "Left" || joinKind == "LeftSemi" || joinKind == "LeftOnly");
 }
 
 // For row storage read we create a separate stage.
@@ -85,7 +89,7 @@ bool TAssignStagesRule::MatchAndApply(TIntrusivePtr<IOperator>& input, TRBOConte
         const auto newStageId = props.StageGraph.AddStage();
         join->Props.StageId = newStageId;
 
-        MaybeSetJoinAlgo(join->Props, ctx);
+        FinalizeJoinPhysicalProps(*join, ctx);
 
         // For cross-join or map join we build a stage with map and broadcast connections
         // FIXME: We assume that right side is small one, map join also can work with hash shuffle connections.
diff --git a/ydb/core/kqp/ut/rbo/kqp_rbo_yql_ut.cpp b/ydb/core/kqp/ut/rbo/kqp_rbo_yql_ut.cpp
index 8efa0809eca..6cadf6f4347 100644
--- a/ydb/core/kqp/ut/rbo/kqp_rbo_yql_ut.cpp
+++ b/ydb/core/kqp/ut/rbo/kqp_rbo_yql_ut.cpp
@@ -1107,6 +1107,8 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
             {filter}
         );
 
+        join.Props.JoinAlgo = NKqp::EJoinAlgoType::GraceJoin;
+        join.Props.UseBlockHashJoin = false;
         const auto joinJson = join.ToJson(0);
         const auto condition = GetStringField(joinJson, "Condition");
         UNIT_ASSERT_C(condition.Contains("t1.id = t2.id"), condition);
@@ -2701,6 +2703,15 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
             auto ast = *result.GetStats()->GetAst();
             const auto isCrossJoin = query.find("AnsiImplicitCrossJoin") != std::string::npos;
             UNIT_ASSERT_C(isCrossJoin || ast.find(joinAlgo) != std::string::npos, TStringBuilder() << "Wrong join algo. Expected: " << joinAlgo);
+            if (!isCrossJoin) {
+                const auto plan = TString{*result.GetStats()->GetPlan()};
+                const auto simplifiedPlan = GetSimplifiedPlan(plan);
+                const TString explainJoinAlgo = useBlockHashJoin ? "BlockHash" : "Grace";
+                const auto* joinOp = FindOperatorByStringField(simplifiedPlan, "JoinAlgo", explainJoinAlgo);
+                UNIT_ASSERT_C(joinOp, plan);
+                const TString explainJoinName = TStringBuilder() << "Join (" << explainJoinAlgo << ")";
+                UNIT_ASSERT_C(GetStringField(*joinOp, "Name").Contains(explainJoinName), plan);
+            }
 
             result =
                 session.ExecuteQuery(query, NYdb::NQuery::TTxControl::NoTx(), NYdb::NQuery::TExecuteQuerySettings().ExecMode(NQuery::EExecMode::Execute))
@@ -7926,7 +7937,7 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
         return false;
     }
 
-    bool IsGraceJoinPlanNode(const NJson::TJsonValue& planNode) {
+    bool IsJoinPlanNode(const NJson::TJsonValue& planNode) {
         if (!planNode.IsMap()) {
             return false;
         }
@@ -7935,11 +7946,7 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
         if (auto operators = planMap.find("Operators"); operators != planMap.end()) {
             for (const auto& opNode : operators->second.GetArraySafe()) {
                 const auto& op = opNode.GetMapSafe();
-                const auto opName = op.at("Name").GetStringSafe();
-                const bool isJoin = opName.Contains("Join");
-                const bool isGrace = opName.Contains("Grace") ||
-                    (op.contains("JoinAlgo") && op.at("JoinAlgo").GetStringSafe() == "Grace");
-                if (isJoin && isGrace) {
+                if (op.contains("JoinKind")) {
                     return true;
                 }
             }
@@ -7948,36 +7955,36 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
         return false;
     }
 
-    ui32 CountGraceJoinPlanNodes(const NJson::TJsonValue& planNode) {
+    ui32 CountJoinPlanNodes(const NJson::TJsonValue& planNode) {
         if (!planNode.IsMap()) {
             return 0;
         }
 
-        ui32 count = IsGraceJoinPlanNode(planNode) ? 1 : 0;
+        ui32 count = IsJoinPlanNode(planNode) ? 1 : 0;
 
         const auto& planMap = planNode.GetMapSafe();
         if (auto plans = planMap.find("Plans"); plans != planMap.end()) {
             for (const auto& child : plans->second.GetArraySafe()) {
-                count += CountGraceJoinPlanNodes(child);
+                count += CountJoinPlanNodes(child);
             }
         }
 
         return count;
     }
 
-    ui32 CountGraceJoinPlanNodes(const TString& plan) {
+    ui32 CountJoinPlanNodes(const TString& plan) {
         NJson::TJsonValue planRoot;
         NJson::ReadJsonTree(plan, &planRoot, true);
-        return CountGraceJoinPlanNodes(planRoot.GetMapSafe().at("SimplifiedPlan"));
+        return CountJoinPlanNodes(planRoot.GetMapSafe().at("SimplifiedPlan"));
     }
 
-    bool HasGraceJoinWithBothInputsHashShuffled(const NJson::TJsonValue& planNode) {
+    bool HasJoinWithBothInputsHashShuffled(const NJson::TJsonValue& planNode) {
         if (!planNode.IsMap()) {
             return false;
         }
 
         const auto& planMap = planNode.GetMapSafe();
-        if (IsGraceJoinPlanNode(planNode)) {
+        if (IsJoinPlanNode(planNode)) {
             if (auto plans = planMap.find("Plans"); plans != planMap.end()) {
                 const auto& children = plans->second.GetArraySafe();
                 if (children.size() == 2 && IsHashShufflePlanNode(children[0]) && IsHashShufflePlanNode(children[1])) {
@@ -7988,7 +7995,7 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
 
         if (auto plans = planMap.find("Plans"); plans != planMap.end()) {
             for (const auto& child : plans->second.GetArraySafe()) {
-                if (HasGraceJoinWithBothInputsHashShuffled(child)) {
+                if (HasJoinWithBothInputsHashShuffled(child)) {
                     return true;
                 }
             }
@@ -7997,10 +8004,10 @@ Y_UNIT_TEST_SUITE(KqpRboYql) {
         return false;
     }
 
-    bool HasGraceJoinWithBothInputsHashShuffled(const TString& plan) {
+    bool HasJoinWithBothInputsHashShuffled(const TString& plan) {
         NJson::TJsonValue planRoot;
         NJson::ReadJsonTree(plan, &planRoot, true);
-        return HasGraceJoinWithBothInputsHashShuffled(planRoot.GetMapSafe().at("SimplifiedPlan"));
+        return HasJoinWithBothInputsHashShuffled(planRoot.GetMapSafe().at("SimplifiedPlan"));
     }
 
     NKikimrKqp::TKqpSetting MakeHashCompatibilityStatsSetting(const TVector<TString>& tables) {
@@ -8290,9 +8297,12 @@ PRAGMA ydb.OptimizerHints = '
         NYdb::NConsoleClient::TQueryPlanPrinter queryPlanPrinter(NYdb::NConsoleClient::EDataFormat::PrettyTable, true, Cout, 0);
         queryPlanPrinter.Print(plan);
 
+        const auto simplifiedPlan = GetSimplifiedPlan(plan);
+        UNIT_ASSERT_C(FindOperatorByStringField(simplifiedPlan, "JoinAlgo", "BlockHash"), plan);
+
         const auto hashShuffles = CollectHashShuffleDescriptions(plan);
 
-        UNIT_ASSERT_VALUES_EQUAL_C(CountGraceJoinPlanNodes(plan), 1u, plan);
+        UNIT_ASSERT_VALUES_EQUAL_C(CountJoinPlanNodes(plan), 1u, plan);
         UNIT_ASSERT_VALUES_EQUAL_C(hashShuffles.size(), 2u, plan);
 
         const bool hasCustomerShuffle = std::any_of(
@@ -8309,8 +8319,8 @@ PRAGMA ydb.OptimizerHints = '
             });
 
         UNIT_ASSERT_C(
-            HasGraceJoinWithBothInputsHashShuffled(plan),
-            TStringBuilder() << "Expected a GraceJoin with HashShuffle on both inputs, got: "
+            HasJoinWithBothInputsHashShuffled(plan),
+            TStringBuilder() << "Expected a join with HashShuffle on both inputs, got: "
                              << JoinSeq(", ", hashShuffles) << "\n" << plan);
         UNIT_ASSERT_C(
             hasCustomerShuffle && hasOrdersShuffle,