Add CSV delimter support in CSVRowOutputStream

This commit is contained in:
Ivan Zhukov 2018-04-22 01:52:24 +03:00
parent 91cb03bdaf
commit a138ab2820
3 changed files with 17 additions and 10 deletions

View File

@ -7,8 +7,8 @@ namespace DB
{
CSVRowOutputStream::CSVRowOutputStream(WriteBuffer & ostr_, const Block & sample_, bool with_names_, bool with_types_)
: ostr(ostr_), sample(sample_), with_names(with_names_), with_types(with_types_)
CSVRowOutputStream::CSVRowOutputStream(WriteBuffer & ostr_, const Block & sample_, const char delimiter_, bool with_names_, bool with_types_)
: ostr(ostr_), sample(sample_), delimiter(delimiter_), with_names(with_names_), with_types(with_types_)
{
size_t columns = sample.columns();
data_types.resize(columns);
@ -32,7 +32,7 @@ void CSVRowOutputStream::writePrefix()
for (size_t i = 0; i < columns; ++i)
{
writeCSVString(sample.safeGetByPosition(i).name, ostr);
writeChar(i == columns - 1 ? '\n' : ',', ostr);
writeChar(i == columns - 1 ? '\n' : delimiter, ostr);
}
}
@ -41,7 +41,7 @@ void CSVRowOutputStream::writePrefix()
for (size_t i = 0; i < columns; ++i)
{
writeCSVString(sample.safeGetByPosition(i).type->getName(), ostr);
writeChar(i == columns - 1 ? '\n' : ',', ostr);
writeChar(i == columns - 1 ? '\n' : delimiter, ostr);
}
}
}
@ -55,7 +55,7 @@ void CSVRowOutputStream::writeField(const IColumn & column, const IDataType & ty
void CSVRowOutputStream::writeFieldDelimiter()
{
writeChar(',', ostr);
writeChar(delimiter, ostr);
}

View File

@ -19,7 +19,7 @@ public:
/** with_names - output in the first line a header with column names
* with_types - output in the next line header with the names of the types
*/
CSVRowOutputStream(WriteBuffer & ostr_, const Block & sample_, bool with_names_ = false, bool with_types_ = false);
CSVRowOutputStream(WriteBuffer & ostr_, const Block & sample_, const char delimiter_, bool with_names_ = false, bool with_types_ = false);
void writeField(const IColumn & column, const IDataType & type, size_t row_num) override;
void writeFieldDelimiter() override;
@ -44,6 +44,7 @@ protected:
WriteBuffer & ostr;
const Block sample;
const char delimiter;
bool with_names;
bool with_types;
DataTypes data_types;

View File

@ -153,10 +153,16 @@ static BlockOutputStreamPtr getOutputImpl(const String & name, WriteBuffer & buf
return std::make_shared<BlockOutputStreamFromRowOutputStream>(std::make_shared<TabSeparatedRowOutputStream>(buf, sample, true, true), sample);
else if (name == "TabSeparatedRaw" || name == "TSVRaw")
return std::make_shared<BlockOutputStreamFromRowOutputStream>(std::make_shared<TabSeparatedRawRowOutputStream>(buf, sample), sample);
else if (name == "CSV")
return std::make_shared<BlockOutputStreamFromRowOutputStream>(std::make_shared<CSVRowOutputStream>(buf, sample), sample);
else if (name == "CSVWithNames")
return std::make_shared<BlockOutputStreamFromRowOutputStream>(std::make_shared<CSVRowOutputStream>(buf, sample, true), sample);
else if (name == "CSV" || name == "CSVWithNames")
{
// TODO: remove self-repeating
String csv_delimiter = settings.format_csv_delimiter.toString();
if (csv_delimiter.size() != 1)
throw Exception("A format_csv_delimiter setting has to be an exactly one character long");
bool with_names = name == "CSVWithNames";
return std::make_shared<BlockOutputStreamFromRowOutputStream>(std::make_shared<CSVRowOutputStream>(buf, sample, csv_delimiter[0], with_names), sample);
}
else if (name == "Pretty")
return std::make_shared<PrettyBlockOutputStream>(buf, sample, false, settings.output_format_pretty_max_rows, context);
else if (name == "PrettyCompact")