Skip to content

Commit

Permalink
Pretty print results in benchmarks (ydb-platform#11050)
Browse files Browse the repository at this point in the history
  • Loading branch information
@iddqdex
iddqdex authored Oct 29, 2024
1 parent c00aee2 commit cd2ba06
Show file tree
Hide file tree
Showing 7 changed files with 84 additions and 83 deletions.
77 changes: 33 additions & 44 deletions ydb/public/lib/ydb_cli/commands/benchmark_utils.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <library/cpp/colorizer/colors.h>
#include <library/cpp/json/json_writer.h>
#include <library/cpp/string_utils/csv/csv.h>
#include <library/cpp/digest/md5/md5.h>

#include <ydb/public/sdk/cpp/client/ydb_table/table.h>
#include <ydb/public/lib/ydb_cli/common/pretty_table.h>
Expand Down Expand Up @@ -160,9 +161,12 @@ class IQueryResultScanner {
YDB_READONLY_DEF(TString, QueryPlan);
YDB_READONLY_DEF(TString, PlanAst);
YDB_ACCESSOR_DEF(TString, DeadlineName);

TQueryBenchmarkResult::TRawResults RawResults;
public:
virtual ~IQueryResultScanner() = default;
TQueryBenchmarkResult::TRawResults&& ExtractRawResults() {
return std::move(RawResults);
}
virtual void OnStart(const TVector<NYdb::TColumn>& columns) = 0;
virtual void OnBeforeRow() = 0;
virtual void OnAfterRow() = 0;
Expand Down Expand Up @@ -208,11 +212,11 @@ class IQueryResultScanner {
}

if (streamPart.HasResultSet()) {
auto result = streamPart.ExtractResultSet();
auto columns = result.GetColumnsMeta();
RawResults.emplace_back(streamPart.ExtractResultSet());
auto columns = RawResults.back().GetColumnsMeta();

OnStart(columns);
NYdb::TResultSetParser parser(result);
NYdb::TResultSetParser parser(RawResults.back());
while (parser.TryNextRow()) {
OnBeforeRow();
for (ui32 i = 0; i < columns.size(); ++i) {
Expand Down Expand Up @@ -262,37 +266,6 @@ class TQueryResultScannerComposite: public IQueryResultScanner {
}
};

class TYSONResultScanner: public IQueryResultScanner {
private:
TStringStream ResultString;
mutable std::unique_ptr<NYson::TYsonWriter> Writer;
public:
TYSONResultScanner() {
}
TString GetResult() const {
Writer.reset();
return ResultString.Str();
}
virtual void OnStart(const TVector<NYdb::TColumn>& /*columns*/) override {
Writer = std::make_unique<NYson::TYsonWriter>(&ResultString, NYson::EYsonFormat::Text, ::NYson::EYsonType::Node, true);
Writer->OnBeginList();
}
virtual void OnBeforeRow() override {
Writer->OnListItem();
Writer->OnBeginList();
}
virtual void OnAfterRow() override {
Writer->OnEndList();
}
virtual void OnRowItem(const NYdb::TColumn& /*c*/, const NYdb::TValue& value) override {
Writer->OnListItem();
FormatValueYson(value, *Writer);
}
virtual void OnFinish() override {
Writer->OnEndList();
}
};

class TCSVResultScanner: public IQueryResultScanner, public TQueryResultInfo {
public:
TCSVResultScanner() {
Expand Down Expand Up @@ -328,26 +301,24 @@ TQueryBenchmarkResult ExecuteImpl(const TString& query, NTable::TTableClient& cl
TStreamExecScanQuerySettings settings;
settings.CollectQueryStats(ECollectQueryStatsMode::Full);
settings.Explain(explainOnly);
if (auto error = SetTimeoutSettings(settings, deadline)) {
if (const auto error = SetTimeoutSettings(settings, deadline)) {
return *error;
}
auto it = client.StreamExecuteScanQuery(query, settings).GetValueSync();
if (auto error = ResultByStatus(it, deadline.Name)) {
if (const auto error = ResultByStatus(it, deadline.Name)) {
return *error;
}

std::shared_ptr<TYSONResultScanner> scannerYson = std::make_shared<TYSONResultScanner>();
std::shared_ptr<TCSVResultScanner> scannerCSV = std::make_shared<TCSVResultScanner>();
TQueryResultScannerComposite composite;
composite.SetDeadlineName(deadline.Name);
composite.AddScanner(scannerYson);
composite.AddScanner(scannerCSV);
if (!composite.Scan(it)) {
return TQueryBenchmarkResult::Error(
composite.GetErrorInfo(), composite.GetQueryPlan(), composite.GetPlanAst());
} else {
return TQueryBenchmarkResult::Result(
scannerYson->GetResult(),
composite.ExtractRawResults(),
*scannerCSV,
composite.GetServerTiming(),
composite.GetQueryPlan(),
Expand Down Expand Up @@ -379,18 +350,16 @@ TQueryBenchmarkResult ExecuteImpl(const TString& query, NQuery::TQueryClient& cl
return *error;
}

std::shared_ptr<TYSONResultScanner> scannerYson = std::make_shared<TYSONResultScanner>();
std::shared_ptr<TCSVResultScanner> scannerCSV = std::make_shared<TCSVResultScanner>();
TQueryResultScannerComposite composite;
composite.SetDeadlineName(deadline.Name);
composite.AddScanner(scannerYson);
composite.AddScanner(scannerCSV);
if (!composite.Scan(it)) {
return TQueryBenchmarkResult::Error(
composite.GetErrorInfo(), composite.GetQueryPlan(), composite.GetPlanAst());
} else {
return TQueryBenchmarkResult::Result(
scannerYson->GetResult(),
composite.ExtractRawResults(),
*scannerCSV,
composite.GetServerTiming(),
composite.GetQueryPlan(),
Expand Down Expand Up @@ -629,6 +598,26 @@ bool CompareValue(const NYdb::TValue& v, TStringBuf vExpected) {
}
}

TQueryResultInfo::TColumnsRemap TQueryResultInfo::GetColumnsRemap() const {
TColumnsRemap result;
ui32 idx = 0;
for (auto&& i : Columns) {
result.emplace(i.Name, idx++);
}
return result;
}

TString TQueryResultInfo::CalcHash() const {
MD5 hasher;
for (const auto& row: Result) {
for (const auto& v: row) {
hasher.Update(FormatValueYson(v, NYson::EYsonFormat::Binary));
}
}
char buf[25];
return hasher.End_b64(buf);
}

bool TQueryResultInfo::IsExpected(std::string_view expected) const {
if (expected.empty()) {
return true;
Expand All @@ -647,7 +636,7 @@ bool TQueryResultInfo::IsExpected(std::string_view expected) const {

std::vector<ui32> columnIndexes;
{
const std::map<TString, ui32> columns = GetColumnsRemap();
const auto columns = GetColumnsRemap();
auto copy = expectedLines.front();
NCsvFormat::CsvSplitter splitter(copy);
while (true) {
Expand Down
34 changes: 14 additions & 20 deletions ydb/public/lib/ydb_cli/commands/benchmark_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,56 +31,50 @@ struct TTestInfo {
};

class TQueryResultInfo {
YDB_READONLY_PROTECT_DEF(std::vector<std::vector<NYdb::TValue>>, Result);
YDB_READONLY_PROTECT_DEF(TVector<NYdb::TColumn>, Columns);
protected:
std::vector<std::vector<NYdb::TValue>> Result;
TVector<NYdb::TColumn> Columns;
public:
std::map<TString, ui32> GetColumnsRemap() const {
std::map<TString, ui32> result;
ui32 idx = 0;
for (auto&& i : Columns) {
result.emplace(i.Name, idx++);
}
return result;
}
using TColumnsRemap = std::map<TString, ui32>;
TColumnsRemap GetColumnsRemap() const;

const std::vector<std::vector<NYdb::TValue>>& GetResult() const {
return Result;
}
const TVector<NYdb::TColumn>& GetColumns() const {
return Columns;
}
public:
bool IsExpected(std::string_view expected) const;
TString CalcHash() const;
};

class TQueryBenchmarkResult {
public:
using TRawResults = TVector<NYdb::TResultSet>;

private:
YDB_READONLY_DEF(TString, ErrorInfo);
YDB_READONLY_DEF(TString, YSONResult);
YDB_READONLY_DEF(TRawResults, RawResults);
YDB_READONLY_DEF(TQueryResultInfo, QueryResult);
YDB_READONLY_DEF(TDuration, ServerTiming);
YDB_READONLY_DEF(TString, QueryPlan);
YDB_READONLY_DEF(TString, PlanAst);
TQueryBenchmarkResult() = default;
public:
static TQueryBenchmarkResult Result(const TString& yson, const TQueryResultInfo& queryResult,
static TQueryBenchmarkResult Result(TRawResults&& rawResults, const TQueryResultInfo& queryResult,
const TDuration& serverTiming, const TString& queryPlan, const TString& planAst)
{
TQueryBenchmarkResult result;
result.YSONResult = yson;
result.RawResults = std::move(rawResults);
result.QueryResult = queryResult;
result.ServerTiming = serverTiming;
result.QueryPlan = queryPlan;
result.PlanAst = planAst;
return result;
}

static TQueryBenchmarkResult Error(const TString& error, const TString& queryPlan, const TString& planAst) {
TQueryBenchmarkResult result;
result.ErrorInfo = error;
result.QueryPlan = queryPlan;
result.PlanAst = planAst;
return result;
}

operator bool() const {
return !ErrorInfo;
}
Expand Down
22 changes: 17 additions & 5 deletions ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,17 +366,20 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
serverTimings.emplace_back(res.GetServerTiming());
++successIteration;
if (successIteration == 1) {
outFStream << queryN << ": " << Endl
<< res.GetYSONResult() << Endl << Endl;
outFStream << queryN << ": " << Endl;
PrintResult(res, outFStream);
}
if ((!prevResult || *prevResult != res.GetYSONResult()) && !res.GetQueryResult().IsExpected(qInfo.ExpectedResult)) {
const auto resHash = res.GetQueryResult().CalcHash();
if ((!prevResult || *prevResult != resHash) && !res.GetQueryResult().IsExpected(qInfo.ExpectedResult)) {
outFStream << queryN << ":" << Endl <<
"Query text:" << Endl <<
query << Endl << Endl <<
"UNEXPECTED DIFF: " << Endl
<< "RESULT: " << Endl << res.GetYSONResult() << Endl
<< "RESULT: " << Endl;
PrintResult(res, outFStream);
outFStream << Endl
<< "EXPECTATION: " << Endl << qInfo.ExpectedResult << Endl;
prevResult = res.GetYSONResult();
prevResult = resHash;
++diffsCount;
}
} else {
Expand Down Expand Up @@ -456,6 +459,15 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
return !someFailQueries;
}

void TWorkloadCommandBenchmark::PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out) const {
TResultSetPrinter printer(EDataFormat::Pretty, []() { return false; }, out, 120);
for(const auto& r: res.GetRawResults()) {
printer.Print(r);
printer.Reset();
}
out << Endl << Endl;
}

void TWorkloadCommandBenchmark::SavePlans(const BenchmarkUtils::TQueryBenchmarkResult& res, ui32 queryNum, const TStringBuf name) const {
if (!PlanFileName) {
return;
Expand Down
1 change: 1 addition & 0 deletions ydb/public/lib/ydb_cli/commands/ydb_benchmark.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ class TWorkloadCommandBenchmark final: public TWorkloadCommandBase {
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;
BenchmarkUtils::TQueryBenchmarkDeadline GetDeadline() const;

private:
Expand Down
25 changes: 14 additions & 11 deletions ydb/public/lib/ydb_cli/common/format.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -650,10 +650,12 @@ TString TQueryPlanPrinter::JsonToString(const NJson::TJsonValue& jsonValue) {
return jsonValue.GetStringRobust();
}

TResultSetPrinter::TResultSetPrinter(EDataFormat format, std::function<bool()> isInterrupted)
TResultSetPrinter::TResultSetPrinter(EDataFormat format, std::function<bool()> isInterrupted, IOutputStream& output, size_t maxWidth)
: Format(format)
, IsInterrupted(isInterrupted)
, ParquetPrinter(std::make_unique<TResultSetParquetPrinter>(""))
, Output(output)
, MaxWidth(maxWidth)
{}

TResultSetPrinter::~TResultSetPrinter() {
Expand All @@ -677,13 +679,13 @@ void TResultSetPrinter::Print(const TResultSet& resultSet) {
PrintJsonArray(resultSet, EBinaryStringEncoding::Unicode);
break;
case EDataFormat::JsonUnicode:
FormatResultSetJson(resultSet, &Cout, EBinaryStringEncoding::Unicode);
FormatResultSetJson(resultSet, &Output, EBinaryStringEncoding::Unicode);
break;
case EDataFormat::JsonBase64Array:
PrintJsonArray(resultSet, EBinaryStringEncoding::Base64);
break;
case EDataFormat::JsonBase64:
FormatResultSetJson(resultSet, &Cout, EBinaryStringEncoding::Base64);
FormatResultSetJson(resultSet, &Output, EBinaryStringEncoding::Base64);
break;
case EDataFormat::Csv:
PrintCsv(resultSet, ",");
Expand Down Expand Up @@ -714,7 +716,7 @@ void TResultSetPrinter::BeginResultSet() {
switch (Format) {
case EDataFormat::JsonUnicodeArray:
case EDataFormat::JsonBase64Array:
Cout << '[';
Output << '[';
break;
default:
break;
Expand All @@ -725,7 +727,7 @@ void TResultSetPrinter::EndResultSet() {
switch (Format) {
case EDataFormat::JsonUnicodeArray:
case EDataFormat::JsonBase64Array:
Cout << ']' << Endl;
Output << ']' << Endl;
break;
case EDataFormat::Parquet:
ParquetPrinter->Reset();
Expand All @@ -739,7 +741,7 @@ void TResultSetPrinter::EndLineBeforeNextResult() {
switch (Format) {
case EDataFormat::JsonUnicodeArray:
case EDataFormat::JsonBase64Array:
Cout << ',' << Endl;
Output << ',' << Endl;
break;
default:
break;
Expand All @@ -755,6 +757,7 @@ void TResultSetPrinter::PrintPretty(const TResultSet& resultSet) {
}

TPrettyTableConfig tableConfig;
tableConfig.MaxWidth(MaxWidth);
if (!FirstPart) {
tableConfig.WithoutHeader();
}
Expand All @@ -767,7 +770,7 @@ void TResultSetPrinter::PrintPretty(const TResultSet& resultSet) {
}
}

Cout << table;
Output << table;
}

void TResultSetPrinter::PrintJsonArray(const TResultSet& resultSet, EBinaryStringEncoding encoding) {
Expand All @@ -782,7 +785,7 @@ void TResultSetPrinter::PrintJsonArray(const TResultSet& resultSet, EBinaryStrin
if (firstRow) {
firstRow = false;
}
NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, &Cout);
NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, &Output);
FormatResultRowJson(parser, columns, writer, encoding);
}
}
Expand All @@ -792,12 +795,12 @@ void TResultSetPrinter::PrintCsv(const TResultSet& resultSet, const char* delim)
TResultSetParser parser(resultSet);
while (parser.TryNextRow()) {
for (ui32 i = 0; i < columns.size(); ++i) {
Cout << FormatValueJson(parser.GetValue(i), EBinaryStringEncoding::Unicode);
Output << FormatValueJson(parser.GetValue(i), EBinaryStringEncoding::Unicode);
if (i < columns.size() - 1) {
Cout << delim;
Output << delim;
}
}
Cout << Endl;
Output << Endl;
}
}

Expand Down
4 changes: 3 additions & 1 deletion ydb/public/lib/ydb_cli/common/format.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ class TCommandWithMessagingFormat {

class TResultSetPrinter {
public:
TResultSetPrinter(EDataFormat format, std::function<bool()> isInterrupted = []() { return false; });
TResultSetPrinter(EDataFormat format, std::function<bool()> isInterrupted = []() { return false; }, IOutputStream& output = Cout, size_t maxWidth = 0);
~TResultSetPrinter();

void Print(const TResultSet& resultSet);
Expand All @@ -126,6 +126,8 @@ class TResultSetPrinter {
EDataFormat Format;
std::function<bool()> IsInterrupted;
std::unique_ptr<TResultSetParquetPrinter> ParquetPrinter;
IOutputStream& Output;
size_t MaxWidth;
};

class TQueryPlanPrinter {
Expand Down
Loading

0 comments on commit cd2ba06

Please sign in to comment.