[KQP] Shuffle elimination flag bug fix (#16367)

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/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -545,7 +545,8 @@ void BuildKqpStageChannels(TKqpTasksGraph& tasksGraph, TStageInfo& stageInfo,
             columnShardHashV1Params.SourceTableKeyColumnTypes->reserve(columnShardHashV1.KeyColumnTypesSize());
             for (const auto& keyColumnType: columnShardHashV1.GetKeyColumnTypes()) {
                 auto typeId = static_cast<NScheme::TTypeId>(keyColumnType);
-                auto typeInfo = NScheme::TTypeInfo{typeId};
+                auto typeInfo =
+                    typeId == NScheme::NTypeIds::Decimal? NScheme::TTypeInfo(NKikimr::NScheme::TDecimalType::Default()): NScheme::TTypeInfo(typeId);
                 columnShardHashV1Params.SourceTableKeyColumnTypes->push_back(typeInfo);
             }
             break;

ydb/core/kqp/opt/logical/kqp_opt_log.cpp

@@ -163,7 +163,7 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
     TMaybeNode<TExprBase> OptimizeEquiJoinWithCosts(TExprBase node, TExprContext& ctx) {
         auto maxDPhypDPTableSize = Config->MaxDPHypDPTableSize.Get().GetOrElse(TDqSettings::TDefault::MaxDPHypDPTableSize);
         auto optLevel = Config->CostBasedOptimizationLevel.Get().GetOrElse(Config->DefaultCostBasedOptimizationLevel);
-        bool enableShuffleElimination = KqpCtx.Config->OptShuffleElimination.Get().GetOrElse(Config->DefaultEnableShuffleElimination);
+        bool enableShuffleElimination = KqpCtx.Config->OptShuffleElimination.Get().GetOrElse(KqpCtx.Config->DefaultEnableShuffleElimination);
         auto providerCtx = TKqpProviderContext(KqpCtx, optLevel);
         auto opt = std::unique_ptr<IOptimizerNew>(MakeNativeOptimizerNew(providerCtx, maxDPhypDPTableSize, ctx, enableShuffleElimination));
         TExprBase output = DqOptimizeEquiJoinWithCosts(node, ctx, TypesCtx, optLevel,

ydb/core/kqp/opt/physical/kqp_opt_phy.cpp

@@ -452,7 +452,7 @@ class TKqpPhysicalOptTransformer : public TOptimizeTransformerBase {
         bool shuffleEliminationWithMap = KqpCtx.Config->OptShuffleEliminationWithMap.Get().GetOrElse(true);
         bool rightCollectStage = !KqpCtx.Config->AllowMultiBroadcasts;
         TExprBase output = DqBuildJoin(node, ctx, optCtx, *getParents(), IsGlobal,
-            pushLeftStage, KqpCtx.Config->GetHashJoinMode(), false, KqpCtx.Config->UseGraceJoinCoreForMap.Get().GetOrElse(false), KqpCtx.Config->OptShuffleElimination.Get().GetOrElse(false), shuffleEliminationWithMap,
+            pushLeftStage, KqpCtx.Config->GetHashJoinMode(), false, KqpCtx.Config->UseGraceJoinCoreForMap.Get().GetOrElse(false), KqpCtx.Config->OptShuffleElimination.Get().GetOrElse(KqpCtx.Config->DefaultEnableShuffleElimination), shuffleEliminationWithMap,
             rightCollectStage
         );
         DumpAppliedRule("BuildJoin", node.Ptr(), output.Ptr(), ctx);

ydb/core/kqp/query_compiler/kqp_query_compiler.cpp

@@ -919,7 +919,7 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             i.MutableProgram()->MutableSettings()->SetLevelDataPrediction(rPredictor.GetLevelDataVolume(i.GetProgram().GetSettings().GetStageLevel()));
         }
 
-        txProto.SetEnableShuffleElimination(Config->OptShuffleElimination.Get().GetOrElse(false));
+        txProto.SetEnableShuffleElimination(Config->OptShuffleElimination.Get().GetOrElse(Config->DefaultEnableShuffleElimination));
         txProto.SetHasEffects(hasEffectStage);
 
         for (const auto& paramBinding : tx.ParamBindings()) {
@@ -1336,13 +1336,16 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
                 shuffleProto.AddKeyColumns(TString(keyColumn));
             }
 
-            if (Config->OptShuffleElimination.Get().GetOrElse(false)) {
+            if (Config->OptShuffleElimination.Get().GetOrElse(Config->DefaultEnableShuffleElimination)) {
                 auto& columnHashV1 = *shuffleProto.MutableColumnShardHashV1();
 
                 const auto& outputType = NYql::NDq::GetDqConnectionType(connection, ctx);
                 auto structType = outputType->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
                 for (const auto& column: shuffle.KeyColumns().Ptr()->Children()) {
                     auto ty = NYql::NDq::GetColumnType(connection, *structType, column->Content(), column->Pos(), ctx);
+                    if (ty->GetKind() == ETypeAnnotationKind::List) {
+                        ty = ty->Cast<TListExprType>()->GetItemType();
+                    }
                     NYql::NUdf::EDataSlot slot;
                     switch (ty->GetKind()) {
                         case ETypeAnnotationKind::Data: {
@@ -1351,6 +1354,9 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
                         }
                         case ETypeAnnotationKind::Optional: {
                             auto optionalType = ty->Cast<TOptionalExprType>()->GetItemType();
+                            if (optionalType->GetKind() == ETypeAnnotationKind::List) {
+                                optionalType = optionalType->Cast<TListExprType>()->GetItemType();
+                            }
                             Y_ENSURE(
                                 optionalType->GetKind() == ETypeAnnotationKind::Data,
                                 TStringBuilder{} << "Can't retrieve type from optional" << static_cast<std::int64_t>(optionalType->GetKind()) << "for ColumnHashV1 Shuffling"