Merge 2b7a4f6 into 41f47f8

Author: GrigoriyPA

ydb/tests/tools/kqprun/src/actors.cpp

@@ -50,23 +50,29 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
         auto resultSetIndex = ev->Get()->Record.GetQueryResultIndex();
         if (resultSetIndex >= ResultSets_.size()) {
             ResultSets_.resize(resultSetIndex + 1);
+            ResultSetSizes_.resize(resultSetIndex + 1, 0);
         }
 
         if (!ResultSets_[resultSetIndex].truncated()) {
+            ui64& resultSetSize = ResultSetSizes_[resultSetIndex];
             for (auto& row : *ev->Get()->Record.MutableResultSet()->mutable_rows()) {
                 if (static_cast<ui64>(ResultSets_[resultSetIndex].rows_size()) >= ResultRowsLimit_) {
                     ResultSets_[resultSetIndex].set_truncated(true);
                     break;
                 }
 
-                if (ResultSets_[resultSetIndex].ByteSizeLong() + row.ByteSizeLong() > ResultSizeLimit_) {
+                auto rowSize = row.ByteSizeLong();
+                if (resultSetSize + rowSize > ResultSizeLimit_) {
                     ResultSets_[resultSetIndex].set_truncated(true);
                     break;
                 }
 
+                resultSetSize += rowSize;
                 *ResultSets_[resultSetIndex].add_rows() = std::move(row);
             }
-            *ResultSets_[resultSetIndex].mutable_columns() = ev->Get()->Record.GetResultSet().columns();
+            if (!ResultSets_[resultSetIndex].columns_size()) {
+                *ResultSets_[resultSetIndex].mutable_columns() = ev->Get()->Record.GetResultSet().columns();
+            }
         }
 
         Send(ev->Sender, response.Release());
@@ -89,6 +95,7 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     ui64 ResultRowsLimit_;
     ui64 ResultSizeLimit_;
     std::vector<Ydb::ResultSet>& ResultSets_;
+    std::vector<ui64> ResultSetSizes_;
     TProgressCallback ProgressCallback_;
 };