[KQP] Shuffle elimination flag bug fix

Author: pashandor789

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -1652,11 +1652,12 @@ class TKqpExecuterBase : public TActor<TDerived> {
         auto& stage = stageInfo.Meta.GetStage(stageInfo.Id);
 
         auto& columnShardHashV1Params = stageInfo.Meta.ColumnShardHashV1Params;
-        if (enableShuffleElimination && stageInfo.Meta.ColumnTableInfoPtr) {
+        if (enableShuffleElimination && stage.GetIsShuffleEliminated() && stageInfo.Meta.ColumnTableInfoPtr) {
             const auto& tableDesc = stageInfo.Meta.ColumnTableInfoPtr->Description;
             columnShardHashV1Params.SourceShardCount = tableDesc.GetColumnShardCount();
             columnShardHashV1Params.SourceTableKeyColumnTypes = std::make_shared<TVector<NScheme::TTypeInfo>>();
             for (const auto& column: tableDesc.GetSharding().GetHashSharding().GetColumns()) {
+                Y_ENSURE(stageInfo.Meta.TableConstInfo->Columns.contains(column), TStringBuilder{} << "Table doesn't have column: " << column);
                 auto columnType = stageInfo.Meta.TableConstInfo->Columns.at(column).Type;
                 columnShardHashV1Params.SourceTableKeyColumnTypes->push_back(columnType);
             }

ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp

@@ -350,7 +350,7 @@ class TOptimizerNativeNew: public IOptimizerNew {
         if (postEnumerationShuffleElimination) {
             EliminateShuffles(hypergraph, bestJoinOrder, orderingsFSM);
         }
-        auto resTree = ConvertFromInternal(bestJoinOrder, fdStorage);
+        auto resTree = ConvertFromInternal(bestJoinOrder, fdStorage, EnableShuffleElimination);
         AddMissingConditions(hypergraph, resTree);
         return resTree;
     }

ydb/library/yql/dq/opt/dq_opt_join_tree_node.cpp

@@ -24,7 +24,8 @@ std::shared_ptr<TJoinOptimizerNodeInternal> MakeJoinInternal(
 
 std::shared_ptr<TJoinOptimizerNode> ConvertFromInternal(
     const std::shared_ptr<IBaseOptimizerNode>& internal,
-    const TFDStorage& fdStorage
+    const TFDStorage& fdStorage,
+    bool enableShuffleElimination
 ) {
     Y_ENSURE(internal->Kind == EOptimizerNodeKind::JoinNodeType);
 
@@ -38,31 +39,40 @@ std::shared_ptr<TJoinOptimizerNode> ConvertFromInternal(
     auto right = join->RightArg;
 
     if (left->Kind == EOptimizerNodeKind::JoinNodeType) {
-        left = ConvertFromInternal(left, fdStorage);
+        left = ConvertFromInternal(left, fdStorage, enableShuffleElimination);
     }
     if (right->Kind == EOptimizerNodeKind::JoinNodeType) {
-        right = ConvertFromInternal(right, fdStorage);
+        right = ConvertFromInternal(right, fdStorage, enableShuffleElimination);
     }
 
     auto newJoin = std::make_shared<TJoinOptimizerNode>(left, right, join->LeftJoinKeys, join->RightJoinKeys, join->JoinType, join->JoinAlgo, join->LeftAny, join->RightAny);
     newJoin->Stats = std::move(join->Stats);
 
-    
-    if (join->ShuffleLeftSideByOrderingIdx != -1) {
+    if (!enableShuffleElimination && join->JoinAlgo == EJoinAlgoType::GraceJoin) {
+        left->Stats.ShuffledByColumns =
+            TIntrusivePtr<TOptimizerStatistics::TShuffledByColumns>(
+                new TOptimizerStatistics::TShuffledByColumns(join->LeftJoinKeys)
+            );
+    } else if (join->ShuffleLeftSideByOrderingIdx != -1) {
         auto shuffledBy = fdStorage.GetInterestingOrderingsColumnNamesByIdx(join->ShuffleLeftSideByOrderingIdx);
 
-        left->Stats.ShuffledByColumns = 
+        left->Stats.ShuffledByColumns =
             TIntrusivePtr<TOptimizerStatistics::TShuffledByColumns>(
                 new TOptimizerStatistics::TShuffledByColumns(std::move(shuffledBy))
             );
     } else {
         left->Stats.ShuffledByColumns = nullptr;
     }
 
-    if (join->ShuffleRightSideByOrderingIdx != -1) {
+    if (!enableShuffleElimination && join->JoinAlgo == EJoinAlgoType::GraceJoin) {
+        right->Stats.ShuffledByColumns =
+            TIntrusivePtr<TOptimizerStatistics::TShuffledByColumns>(
+                new TOptimizerStatistics::TShuffledByColumns(join->RightJoinKeys)
+            );
+    } else if (join->ShuffleRightSideByOrderingIdx != -1) {
         auto shuffledBy = fdStorage.GetInterestingOrderingsColumnNamesByIdx(join->ShuffleRightSideByOrderingIdx);
 
-        right->Stats.ShuffledByColumns = 
+        right->Stats.ShuffledByColumns =
             TIntrusivePtr<TOptimizerStatistics::TShuffledByColumns>(
                 new TOptimizerStatistics::TShuffledByColumns(std::move(shuffledBy))
             );

ydb/library/yql/dq/opt/dq_opt_join_tree_node.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <yql/essentials/core/cbo/cbo_optimizer_new.h> 
+#include <yql/essentials/core/cbo/cbo_optimizer_new.h>
 
 const TString& ToString(NYql::EJoinKind);
 const TString& ToString(NYql::EJoinAlgoType);
@@ -19,15 +19,15 @@ namespace NYql::NDq {
 */
 struct TJoinOptimizerNodeInternal : public IBaseOptimizerNode {
     TJoinOptimizerNodeInternal(
-        const std::shared_ptr<IBaseOptimizerNode>& left, 
+        const std::shared_ptr<IBaseOptimizerNode>& left,
         const std::shared_ptr<IBaseOptimizerNode>& right,
         const TVector<TJoinColumn>& leftJoinKeys,
-        const TVector<TJoinColumn>& rightJoinKeys, 
-        const EJoinKind joinType, 
+        const TVector<TJoinColumn>& rightJoinKeys,
+        const EJoinKind joinType,
         const EJoinAlgoType joinAlgo,
         const bool leftAny,
         const bool rightAny
-    ) 
+    )
         : IBaseOptimizerNode(JoinNodeType)
         , LeftArg(left)
         , RightArg(right)
@@ -55,9 +55,9 @@ struct TJoinOptimizerNodeInternal : public IBaseOptimizerNode {
         stream << ToString(JoinType) << "," << ToString(JoinAlgo) << " ";
 
         for (size_t i = 0; i < LeftJoinKeys.size(); ++i){
-            stream 
+            stream
                 << LeftJoinKeys[i].RelName << "." << LeftJoinKeys[i].AttributeName
-                << "=" 
+                << "="
                 << RightJoinKeys[i].RelName << "." << RightJoinKeys[i].AttributeName << ",";
         }
         stream << "\n";
@@ -75,7 +75,7 @@ struct TJoinOptimizerNodeInternal : public IBaseOptimizerNode {
             stream << "   ";
         }
         stream << "  ";
-        stream << "Shuffled By: " << ShuffleRightSideByOrderingIdx << "\n"; 
+        stream << "Shuffled By: " << ShuffleRightSideByOrderingIdx << "\n";
         RightArg->Print(stream, ntabs + 1);
     }
 
@@ -90,7 +90,7 @@ struct TJoinOptimizerNodeInternal : public IBaseOptimizerNode {
 
     // for interesting orderings framework
     std::int64_t ShuffleLeftSideByOrderingIdx  = -1;
-    std::int64_t ShuffleRightSideByOrderingIdx = -1; 
+    std::int64_t ShuffleRightSideByOrderingIdx = -1;
 };
 
 /**
@@ -118,7 +118,8 @@ std::shared_ptr<TJoinOptimizerNodeInternal> MakeJoinInternal(
 */
 std::shared_ptr<TJoinOptimizerNode> ConvertFromInternal(
     const std::shared_ptr<IBaseOptimizerNode>& internal,
-    const TFDStorage& fdStorage
+    const TFDStorage& fdStorage,
+    bool enableShuffleElimination
 );
 
 } // namespace NYql::NDq