Truncate query results output in benchmarks (#12263)

5 files changed, 22 insertions, 8 deletions

diff --git a/ydb/apps/ydb/CHANGELOG.md b/ydb/apps/ydb/CHANGELOG.md
index fb01ef94f0..7f8dbe6e0d 100644
--- a/ydb/apps/ydb/CHANGELOG.md
+++ b/ydb/apps/ydb/CHANGELOG.md
@@ -1,3 +1,4 @@
+* Truncate query results output in benchmarks
 
 ## 2.17.0 ##
 
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp b/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp
index 82fed4855c..39b46f3311 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp
@@ -367,7 +367,7 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
                 ++successIteration;
                 if (successIteration == 1) {
                     outFStream << queryN << ": " << Endl;
-                    PrintResult(res, outFStream);
+                    PrintResult(res, outFStream, qInfo.ExpectedResult);
                 }
                 const auto resHash = res.CalcHash();
                 if ((!prevResult || *prevResult != resHash) && !res.IsExpected(qInfo.ExpectedResult)) {
@@ -376,7 +376,7 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
                         query << Endl << Endl <<
                         "UNEXPECTED DIFF: " << Endl
                           << "RESULT: " << Endl;
-                    PrintResult(res, outFStream);
+                    PrintResult(res, outFStream, qInfo.ExpectedResult);
                     outFStream << Endl
                             << "EXPECTATION: " << Endl << qInfo.ExpectedResult << Endl;
                     prevResult = resHash;
@@ -459,9 +459,10 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
     return !someFailQueries;
 }
 
-void TWorkloadCommandBenchmark::PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out) const {
+void TWorkloadCommandBenchmark::PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out, const std::string& expected) const {
     TResultSetPrinter printer(TResultSetPrinter::TSettings()
         .SetOutput(&out)
+        .SetMaxRowsCount(std::max(StringSplitter(expected.c_str()).Split('\n').Count(), (size_t)100))
         .SetFormat(EDataFormat::Pretty).SetMaxWidth(120)
     );
     for (const auto& [i, rr]: res.GetRawResults()) {
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_benchmark.h b/ydb/public/lib/ydb_cli/commands/ydb_benchmark.h
index f331933014..0c17684583 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_benchmark.h
+++ b/ydb/public/lib/ydb_cli/commands/ydb_benchmark.h
@@ -22,7 +22,7 @@ private:
     template <typename TClient>
     bool RunBench(TClient& client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen);
     void SavePlans(const BenchmarkUtils::TQueryBenchmarkResult& res, ui32 queryNum, const TStringBuf name) const;
-    void PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out) const;
+    void PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out, const std::string& expected) const;
     BenchmarkUtils::TQueryBenchmarkDeadline GetDeadline() const;
 
 private:
diff --git a/ydb/public/lib/ydb_cli/common/format.cpp b/ydb/public/lib/ydb_cli/common/format.cpp
index 087bf50444..d5de467386 100644
--- a/ydb/public/lib/ydb_cli/common/format.cpp
+++ b/ydb/public/lib/ydb_cli/common/format.cpp
@@ -766,9 +766,12 @@ void TResultSetPrinter::PrintPretty(const TResultSet& resultSet) {
         tableConfig.WithoutHeader();
     }
     TPrettyTable table(columnNames, tableConfig);
-
-    while (parser.TryNextRow()) {
+    for (size_t printed = 0; parser.TryNextRow(); ++printed) {
         auto& row = table.AddRow();
+        if (Settings.GetMaxRowsCount() && printed >= Settings.GetMaxRowsCount()) {
+            row.FreeText(TStringBuilder() << "And " << (resultSet.RowsCount() - printed) << " more lines, total " << resultSet.RowsCount());
+            break;
+        }
         for (ui32 i = 0; i < columns.size(); ++i) {
             row.Column(i, FormatValueJson(parser.GetValue(i), EBinaryStringEncoding::Unicode));
         }
@@ -782,13 +785,17 @@ void TResultSetPrinter::PrintJsonArray(const TResultSet& resultSet, EBinaryStrin
 
     TResultSetParser parser(resultSet);
     bool firstRow = true;
-    while (parser.TryNextRow()) {
+    for (size_t printed = 0; parser.TryNextRow(); ++printed) {
         if (!firstRow || !FirstPart) {
             EndLineBeforeNextResult();
         }
         if (firstRow) {
             firstRow = false;
         }
+        if (Settings.GetMaxRowsCount() && printed >= Settings.GetMaxRowsCount()) {
+            *Settings.GetOutput() << "And " << (resultSet.RowsCount() - printed) << " more lines, total " << resultSet.RowsCount();
+            break;
+        }
         NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, Settings.GetOutput());
         FormatResultRowJson(parser, columns, writer, encoding);
     }
@@ -806,7 +813,11 @@ void TResultSetPrinter::PrintCsv(const TResultSet& resultSet, const char* delim)
         }
         *Settings.GetOutput() << Endl;
     }
-    while (parser.TryNextRow()) {
+    for (size_t printed = 0; parser.TryNextRow(); ++printed) {
+        if (Settings.GetMaxRowsCount() && printed >= Settings.GetMaxRowsCount()) {
+            *Settings.GetOutput() << "And " << (resultSet.RowsCount() - printed) << " more lines, total " << resultSet.RowsCount() << Endl;
+            break;
+        }
         for (ui32 i = 0; i < columns.size(); ++i) {
             *Settings.GetOutput() << FormatValueJson(parser.GetValue(i), EBinaryStringEncoding::Unicode);
             if (i < columns.size() - 1) {
diff --git a/ydb/public/lib/ydb_cli/common/format.h b/ydb/public/lib/ydb_cli/common/format.h
index 86d65cf947..2a3f8f9254 100644
--- a/ydb/public/lib/ydb_cli/common/format.h
+++ b/ydb/public/lib/ydb_cli/common/format.h
@@ -109,6 +109,7 @@ public:
         YDB_ACCESSOR(EDataFormat, Format, EDataFormat::Pretty);
         YDB_ACCESSOR(std::function<bool()>, IsInterrupted, []() { return false; });
         YDB_ACCESSOR(size_t, MaxWidth, 0);
+        YDB_ACCESSOR(size_t, MaxRowsCount, 0);
         YDB_FLAG_ACCESSOR(CsvWithHeader, false);
         YDB_READONLY(IOutputStream*, Output, &Cout);
     public: