Add more and better checks

Author: abyss7

ydb/core/kqp/opt/kqp_query_plan.cpp

@@ -467,6 +467,10 @@ class TxPlanSerializer {
             if (!TryFromString(wideColumnName, idx)) {
                 return wideColumnName;
             }
+
+            YQL_ENSURE(idx < settings.OutputNarrowType->GetSize(),
+                "Failed to lookup column name for index " << idx << " in type " << settings.OutputNarrowType->ToString());
+
             return TString(settings.OutputNarrowType->GetItems()[idx]->GetName());
         };
 

ydb/core/kqp/opt/peephole/kqp_opt_peephole.cpp

@@ -287,12 +287,6 @@ bool CanPropagateWideBlockThroughChannel(
         return false;
     }
 
-    auto outputItemType = program.Lambda().Ref().GetTypeAnn()->Cast<TStreamExprType>()->GetItemType();
-    if (IsWideBlockType(*outputItemType)) {
-        // output is already wide block
-        return false;
-    }
-
     YQL_ENSURE(stageSettings.OutputNarrowType);
 
     if (!IsCompatibleWithBlocks(program.Pos(), *stageSettings.OutputNarrowType, ctx, typesCtx)) {
@@ -306,6 +300,15 @@ bool CanPropagateWideBlockThroughChannel(
         return false;
     }
 
+    auto typeAnnotation = program.Lambda().Ref().GetTypeAnn();
+
+    YQL_ENSURE(typeAnnotation, "Program for stage " << output.Stage().Ref().UniqueId() << " doesn't have type annotation");
+
+    if (IsWideBlockType(*typeAnnotation->Cast<TStreamExprType>()->GetItemType())) {
+        // output is already wide block
+        return false;
+    }
+
     return true;
 }
 
@@ -445,8 +448,6 @@ TMaybeNode<TKqpPhysicalTx> PeepholeOptimize(const TKqpPhysicalTx& tx, TExprConte
             .ArgsType(ExpandType(stage.Pos(), *ctx.MakeType<TTupleExprType>(argTypes), ctx))
             .Done();
 
-        YQL_ENSURE(programs.emplace(stage.Ref().UniqueId(), program).second);
-
         const bool allowNonDeterministicFunctions = !program.Lambda().Body().Maybe<TKqpEffects>();
 
         TExprNode::TPtr newProgram;
@@ -468,7 +469,7 @@ TMaybeNode<TKqpPhysicalTx> PeepholeOptimize(const TKqpPhysicalTx& tx, TExprConte
         }
 
         optimizedStages.emplace(stage.Ref().UniqueId());
-        programs.at(stage.Ref().UniqueId()) = TKqpProgram(newProgram);
+        YQL_ENSURE(programs.emplace(stage.Ref().UniqueId(), TKqpProgram(newProgram)).second);
     }
 
     TVector<TKqpParamBinding> bindings(tx.ParamBindings().begin(), tx.ParamBindings().end());