ClickHouse/dbms/src/DataStreams/PrettyBlockOutputStream.cpp

#include <sys/ioctl.h>
#include <unistd.h>

#include <DB/Functions/FunctionFactory.h>
#include <DB/Functions/IFunction.h>
#include <DB/DataTypes/DataTypesNumber.h>
#include <DB/DataTypes/DataTypeString.h>
#include <DB/DataStreams/PrettyBlockOutputStream.h>
#include <DB/Columns/ColumnConst.h>
#include <DB/Columns/ColumnsNumber.h>
#include <DB/IO/WriteBuffer.h>
#include <DB/IO/WriteHelpers.h>


namespace DB
{

PrettyBlockOutputStream::PrettyBlockOutputStream(WriteBuffer & ostr_, bool no_escapes_, size_t max_rows_, const Context & context_)
	 : ostr(ostr_), max_rows(max_rows_), no_escapes(no_escapes_), context(context_)
{
	struct winsize w;
	if (0 == ioctl(STDOUT_FILENO, TIOCGWINSZ, &w))
		terminal_width = w.ws_col;
}


void PrettyBlockOutputStream::flush()
{
	ostr.next();
}


void PrettyBlockOutputStream::calculateWidths(Block & block, Widths_t & max_widths, Widths_t & name_widths)
{
	size_t rows = block.rows();
	size_t columns = block.columns();

	max_widths.resize(columns);
	name_widths.resize(columns);

	FunctionPtr visible_width_func = FunctionFactory::instance().get("visibleWidth", context);

	/// Calculate widths of all values.
	for (size_t i = 0; i < columns; ++i)
	{
		ColumnWithTypeAndName column;
		column.type = std::make_shared<DataTypeUInt64>();
		column.name = "visibleWidth(" + block.getByPosition(i).name + ")";

		size_t result_number = block.columns();
		block.insert(column);

		visible_width_func->execute(block, {i}, result_number);
		column.column = block.safeGetByPosition(result_number).column;

		if (const ColumnUInt64 * col = typeid_cast<const ColumnUInt64 *>(&*column.column))
		{
			const ColumnUInt64::Container_t & res = col->getData();
			for (size_t j = 0; j < rows; ++j)
			{
				if (res[j] > max_widths[i])
					max_widths[i] = res[j];
			}
		}
		else if (const ColumnConstUInt64 * col = typeid_cast<const ColumnConstUInt64 *>(&*column.column))
		{
			UInt64 res = col->getData();
			max_widths[i] = res;
		}
		else
			throw Exception("Illegal column " + column.column->getName()
				+ " of result of function " + visible_width_func->getName(),
				ErrorCodes::ILLEGAL_COLUMN);

		/// And also calculate widths for names of columns.
		{
			const String & name = block.getByPosition(i).name;

			Block block_with_name
			{
				{ std::make_shared<ColumnConstString>(1, name), std::make_shared<DataTypeString>(), "name" },
				{ nullptr, std::make_shared<DataTypeUInt64>(), "width" }
			};

			visible_width_func->execute(block_with_name, {0}, 1);

			name_widths[i] = typeid_cast<const ColumnConstUInt64 &>(*block_with_name.getByPosition(1).column).getData();

			if (name_widths[i] > max_widths[i])
				max_widths[i] = name_widths[i];
		}
	}
}


void PrettyBlockOutputStream::write(const Block & block_)
{
	if (total_rows >= max_rows)
	{
		total_rows += block_.rows();
		return;
	}

	/// We will insert here columns with the calculated values of visible lengths.
	Block block = block_;

	size_t rows = block.rows();
	size_t columns = block.columns();

	Widths_t max_widths;
	Widths_t name_widths;
	calculateWidths(block, max_widths, name_widths);

	/// Create separators
	std::stringstream top_separator;
	std::stringstream middle_names_separator;
	std::stringstream middle_values_separator;
	std::stringstream bottom_separator;

	top_separator 			<< "┏";
	middle_names_separator	<< "┡";
	middle_values_separator	<< "├";
	bottom_separator 		<< "└";
	for (size_t i = 0; i < columns; ++i)
	{
		if (i != 0)
		{
			top_separator 			<< "┳";
			middle_names_separator	<< "╇";
			middle_values_separator	<< "┼";
			bottom_separator 		<< "┴";
		}

		for (size_t j = 0; j < max_widths[i] + 2; ++j)
		{
			top_separator 			<< "━";
			middle_names_separator	<< "━";
			middle_values_separator	<< "─";
			bottom_separator 		<< "─";
		}
	}
	top_separator 			<< "┓\n";
	middle_names_separator	<< "┩\n";
	middle_values_separator	<< "┤\n";
	bottom_separator 		<< "┘\n";

	std::string top_separator_s = top_separator.str();
	std::string middle_names_separator_s = middle_names_separator.str();
	std::string middle_values_separator_s = middle_values_separator.str();
	std::string bottom_separator_s = bottom_separator.str();

	/// Output the block
	writeString(top_separator_s, ostr);

	/// Names
	writeCString("┃ ", ostr);
	for (size_t i = 0; i < columns; ++i)
	{
		if (i != 0)
			writeCString(" ┃ ", ostr);

		const ColumnWithTypeAndName & col = block.safeGetByPosition(i);

		if (!no_escapes)
			writeCString("\033[1m", ostr);

		if (col.type->isNumeric())
		{
			for (size_t k = 0; k < max_widths[i] - name_widths[i]; ++k)
				writeChar(' ', ostr);

			writeEscapedString(col.name, ostr);
		}
		else
		{
			writeEscapedString(col.name, ostr);

			for (size_t k = 0; k < max_widths[i] - name_widths[i]; ++k)
				writeChar(' ', ostr);
		}

		if (!no_escapes)
			writeCString("\033[0m", ostr);
	}
	writeCString(" ┃\n", ostr);

	writeString(middle_names_separator_s, ostr);

	for (size_t i = 0; i < rows && total_rows + i < max_rows; ++i)
	{
		if (i != 0)
			writeString(middle_values_separator_s, ostr);

		writeCString("│ ", ostr);

		for (size_t j = 0; j < columns; ++j)
		{
			if (j != 0)
				writeCString(" │ ", ostr);

			const ColumnWithTypeAndName & col = block.safeGetByPosition(j);

			if (col.type->isNumeric())
			{
				size_t width = get<UInt64>((*block.safeGetByPosition(columns + j).column)[i]);
				for (size_t k = 0; k < max_widths[j] - width; ++k)
					writeChar(' ', ostr);

				col.type->serializeTextEscaped(*col.column.get(), i, ostr);
			}
			else
			{
				col.type->serializeTextEscaped(*col.column.get(), i, ostr);

				size_t width = get<UInt64>((*block.safeGetByPosition(columns + j).column)[i]);
				for (size_t k = 0; k < max_widths[j] - width; ++k)
					writeChar(' ', ostr);
			}
		}

		writeCString(" │\n", ostr);
	}

	writeString(bottom_separator_s, ostr);

	total_rows += rows;
}


void PrettyBlockOutputStream::writeSuffix()
{
	if (total_rows >= max_rows)
	{
		writeCString("  Showed first ", ostr);
		writeIntText(max_rows, ostr);
		writeCString(".\n", ostr);
	}

	total_rows = 0;
	writeTotals();
	writeExtremes();
}


void PrettyBlockOutputStream::writeTotals()
{
	if (totals)
	{
		writeCString("\nTotals:\n", ostr);
		write(totals);
	}
}


void PrettyBlockOutputStream::writeExtremes()
{
	if (extremes)
	{
		writeCString("\nExtremes:\n", ostr);
		write(extremes);
	}
}


}