mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-02 20:42:04 +00:00
432 lines
13 KiB
C++
432 lines
13 KiB
C++
#pragma once
|
||
|
||
#include <DB/IO/WriteBufferFromFile.h>
|
||
#include <DB/IO/CompressedWriteBuffer.h>
|
||
#include <DB/IO/HashingWriteBuffer.h>
|
||
|
||
#include <DB/Storages/MergeTree/MergeTreeData.h>
|
||
#include <DB/Common/escapeForFileName.h>
|
||
#include <DB/DataTypes/DataTypeNested.h>
|
||
#include <DB/DataTypes/DataTypeArray.h>
|
||
|
||
|
||
namespace DB
|
||
{
|
||
class IMergedBlockOutputStream : public IBlockOutputStream
|
||
{
|
||
public:
|
||
IMergedBlockOutputStream(
|
||
MergeTreeData & storage_,
|
||
size_t min_compress_block_size_,
|
||
size_t max_compress_block_size_,
|
||
CompressionMethod compression_method_)
|
||
: storage(storage_),
|
||
min_compress_block_size(min_compress_block_size_),
|
||
max_compress_block_size(max_compress_block_size_),
|
||
compression_method(compression_method_)
|
||
{
|
||
}
|
||
|
||
protected:
|
||
using OffsetColumns = std::set<std::string>;
|
||
|
||
struct ColumnStream
|
||
{
|
||
ColumnStream(
|
||
const String & escaped_column_name_,
|
||
const String & data_path,
|
||
const std::string & marks_path,
|
||
size_t max_compress_block_size,
|
||
CompressionMethod compression_method) :
|
||
escaped_column_name(escaped_column_name_),
|
||
plain_file(data_path, max_compress_block_size, O_TRUNC | O_CREAT | O_WRONLY),
|
||
plain_hashing(plain_file), compressed_buf(plain_hashing, compression_method), compressed(compressed_buf),
|
||
marks_file(marks_path, 4096, O_TRUNC | O_CREAT | O_WRONLY), marks(marks_file) {}
|
||
|
||
String escaped_column_name;
|
||
|
||
/// compressed -> compressed_buf -> plain_hashing -> plain_file
|
||
WriteBufferFromFile plain_file;
|
||
HashingWriteBuffer plain_hashing;
|
||
CompressedWriteBuffer compressed_buf;
|
||
HashingWriteBuffer compressed;
|
||
|
||
/// marks -> marks_file
|
||
WriteBufferFromFile marks_file;
|
||
HashingWriteBuffer marks;
|
||
|
||
void finalize()
|
||
{
|
||
compressed.next();
|
||
plain_file.next();
|
||
marks.next();
|
||
}
|
||
|
||
void sync()
|
||
{
|
||
plain_file.sync();
|
||
marks_file.sync();
|
||
}
|
||
|
||
void addToChecksums(MergeTreeData::DataPart::Checksums & checksums, String name = "")
|
||
{
|
||
if (name == "")
|
||
name = escaped_column_name;
|
||
|
||
checksums.files[name + ".bin"].is_compressed = true;
|
||
checksums.files[name + ".bin"].uncompressed_size = compressed.count();
|
||
checksums.files[name + ".bin"].uncompressed_hash = compressed.getHash();
|
||
checksums.files[name + ".bin"].file_size = plain_hashing.count();
|
||
checksums.files[name + ".bin"].file_hash = plain_hashing.getHash();
|
||
|
||
checksums.files[name + ".mrk"].file_size = marks.count();
|
||
checksums.files[name + ".mrk"].file_hash = marks.getHash();
|
||
}
|
||
};
|
||
|
||
using ColumnStreams = std::map<String, std::unique_ptr<ColumnStream>>;
|
||
|
||
void addStream(const String & path, const String & name, const IDataType & type, size_t level = 0, String filename = "")
|
||
{
|
||
String escaped_column_name;
|
||
if (filename.size())
|
||
escaped_column_name = escapeForFileName(filename);
|
||
else
|
||
escaped_column_name = escapeForFileName(name);
|
||
|
||
/// Для массивов используются отдельные потоки для размеров.
|
||
if (const DataTypeArray * type_arr = typeid_cast<const DataTypeArray *>(&type))
|
||
{
|
||
String size_name = DataTypeNested::extractNestedTableName(name)
|
||
+ ARRAY_SIZES_COLUMN_NAME_SUFFIX + toString(level);
|
||
String escaped_size_name = escapeForFileName(DataTypeNested::extractNestedTableName(name))
|
||
+ ARRAY_SIZES_COLUMN_NAME_SUFFIX + toString(level);
|
||
|
||
column_streams[size_name].reset(new ColumnStream(
|
||
escaped_size_name,
|
||
path + escaped_size_name + ".bin",
|
||
path + escaped_size_name + ".mrk",
|
||
max_compress_block_size,
|
||
compression_method));
|
||
|
||
addStream(path, name, *type_arr->getNestedType(), level + 1);
|
||
}
|
||
else
|
||
column_streams[name].reset(new ColumnStream(
|
||
escaped_column_name,
|
||
path + escaped_column_name + ".bin",
|
||
path + escaped_column_name + ".mrk",
|
||
max_compress_block_size,
|
||
compression_method));
|
||
}
|
||
|
||
|
||
/// Записать данные одного столбца.
|
||
void writeData(const String & name, const IDataType & type, const IColumn & column, OffsetColumns & offset_columns, size_t level = 0)
|
||
{
|
||
size_t size = column.size();
|
||
|
||
/// Для массивов требуется сначала сериализовать размеры, а потом значения.
|
||
if (const DataTypeArray * type_arr = typeid_cast<const DataTypeArray *>(&type))
|
||
{
|
||
String size_name = DataTypeNested::extractNestedTableName(name)
|
||
+ ARRAY_SIZES_COLUMN_NAME_SUFFIX + toString(level);
|
||
|
||
if (offset_columns.count(size_name) == 0)
|
||
{
|
||
offset_columns.insert(size_name);
|
||
|
||
ColumnStream & stream = *column_streams[size_name];
|
||
|
||
size_t prev_mark = 0;
|
||
while (prev_mark < size)
|
||
{
|
||
size_t limit = 0;
|
||
|
||
/// Если есть index_offset, то первая засечка идёт не сразу, а после этого количества строк.
|
||
if (prev_mark == 0 && index_offset != 0)
|
||
{
|
||
limit = index_offset;
|
||
}
|
||
else
|
||
{
|
||
limit = storage.index_granularity;
|
||
|
||
/// Уже могло накопиться достаточно данных для сжатия в новый блок.
|
||
if (stream.compressed.offset() >= min_compress_block_size)
|
||
stream.compressed.next();
|
||
|
||
writeIntBinary(stream.plain_hashing.count(), stream.marks);
|
||
writeIntBinary(stream.compressed.offset(), stream.marks);
|
||
}
|
||
|
||
type_arr->serializeOffsets(column, stream.compressed, prev_mark, limit);
|
||
|
||
/// Чтобы вместо засечек, указывающих на конец сжатого блока, были засечки, указывающие на начало следующего.
|
||
stream.compressed.nextIfAtEnd();
|
||
|
||
prev_mark += limit;
|
||
}
|
||
}
|
||
}
|
||
|
||
{
|
||
ColumnStream & stream = *column_streams[name];
|
||
|
||
size_t prev_mark = 0;
|
||
while (prev_mark < size)
|
||
{
|
||
size_t limit = 0;
|
||
|
||
/// Если есть index_offset, то первая засечка идёт не сразу, а после этого количества строк.
|
||
if (prev_mark == 0 && index_offset != 0)
|
||
{
|
||
limit = index_offset;
|
||
}
|
||
else
|
||
{
|
||
limit = storage.index_granularity;
|
||
|
||
/// Уже могло накопиться достаточно данных для сжатия в новый блок.
|
||
if (stream.compressed.offset() >= min_compress_block_size)
|
||
stream.compressed.next();
|
||
|
||
writeIntBinary(stream.plain_hashing.count(), stream.marks);
|
||
writeIntBinary(stream.compressed.offset(), stream.marks);
|
||
}
|
||
|
||
type.serializeBinary(column, stream.compressed, prev_mark, limit);
|
||
|
||
/// Чтобы вместо засечек, указывающих на конец сжатого блока, были засечки, указывающие на начало следующего.
|
||
stream.compressed.nextIfAtEnd();
|
||
|
||
prev_mark += limit;
|
||
}
|
||
}
|
||
}
|
||
|
||
MergeTreeData & storage;
|
||
|
||
ColumnStreams column_streams;
|
||
|
||
/// Смещение до первой строчки блока, для которой надо записать индекс.
|
||
size_t index_offset = 0;
|
||
|
||
size_t min_compress_block_size;
|
||
size_t max_compress_block_size;
|
||
|
||
CompressionMethod compression_method;
|
||
};
|
||
|
||
/** Для записи одного куска. Данные уже отсортированы, относятся к одному месяцу, и пишутся в один кускок.
|
||
*/
|
||
class MergedBlockOutputStream : public IMergedBlockOutputStream
|
||
{
|
||
public:
|
||
MergedBlockOutputStream(
|
||
MergeTreeData & storage_,
|
||
String part_path_,
|
||
const NamesAndTypesList & columns_list_,
|
||
CompressionMethod compression_method)
|
||
: IMergedBlockOutputStream(
|
||
storage_, storage_.context.getSettings().min_compress_block_size,
|
||
storage_.context.getSettings().max_compress_block_size, compression_method),
|
||
columns_list(columns_list_), part_path(part_path_)
|
||
{
|
||
Poco::File(part_path).createDirectories();
|
||
|
||
if (storage.mode != MergeTreeData::Unsorted)
|
||
{
|
||
index_file_stream = new WriteBufferFromFile(part_path + "primary.idx", DBMS_DEFAULT_BUFFER_SIZE, O_TRUNC | O_CREAT | O_WRONLY);
|
||
index_stream = new HashingWriteBuffer(*index_file_stream);
|
||
}
|
||
|
||
for (const auto & it : columns_list)
|
||
addStream(part_path, it.name, *it.type);
|
||
}
|
||
|
||
void write(const Block & block) override
|
||
{
|
||
size_t rows = block.rows();
|
||
|
||
/// Сначала пишем индекс. Индекс содержит значение Primary Key для каждой index_granularity строки.
|
||
typedef std::vector<const ColumnWithNameAndType *> PrimaryColumns;
|
||
PrimaryColumns primary_columns;
|
||
|
||
for (const auto & descr : storage.getSortDescription())
|
||
primary_columns.push_back(
|
||
!descr.column_name.empty()
|
||
? &block.getByName(descr.column_name)
|
||
: &block.getByPosition(descr.column_number));
|
||
|
||
for (size_t i = index_offset; i < rows; i += storage.index_granularity)
|
||
{
|
||
for (PrimaryColumns::const_iterator it = primary_columns.begin(); it != primary_columns.end(); ++it)
|
||
{
|
||
if (storage.mode != MergeTreeData::Unsorted)
|
||
index_vec.push_back((*(*it)->column)[i]);
|
||
|
||
(*it)->type->serializeBinary(index_vec.back(), *index_stream);
|
||
}
|
||
|
||
++marks_count;
|
||
}
|
||
|
||
/// Множество записанных столбцов со смещениями, чтобы не писать общие для вложенных структур столбцы несколько раз
|
||
OffsetColumns offset_columns;
|
||
|
||
/// Теперь пишем данные.
|
||
for (const auto & it : columns_list)
|
||
{
|
||
const ColumnWithNameAndType & column = block.getByName(it.name);
|
||
writeData(column.name, *column.type, *column.column, offset_columns);
|
||
}
|
||
|
||
size_t written_for_last_mark = (storage.index_granularity - index_offset + rows) % storage.index_granularity;
|
||
index_offset = (storage.index_granularity - written_for_last_mark) % storage.index_granularity;
|
||
}
|
||
|
||
void writeSuffix() override
|
||
{
|
||
throw Exception("Method writeSuffix is not supported by MergedBlockOutputStream", ErrorCodes::NOT_IMPLEMENTED);
|
||
}
|
||
|
||
MergeTreeData::DataPart::Checksums writeSuffixAndGetChecksums()
|
||
{
|
||
/// Заканчиваем запись и достаем чексуммы.
|
||
MergeTreeData::DataPart::Checksums checksums;
|
||
|
||
if (storage.mode != MergeTreeData::Unsorted)
|
||
{
|
||
index_stream->next();
|
||
checksums.files["primary.idx"].file_size = index_stream->count();
|
||
checksums.files["primary.idx"].file_hash = index_stream->getHash();
|
||
index_stream = nullptr;
|
||
}
|
||
|
||
for (ColumnStreams::iterator it = column_streams.begin(); it != column_streams.end(); ++it)
|
||
{
|
||
it->second->finalize();
|
||
it->second->addToChecksums(checksums);
|
||
}
|
||
|
||
column_streams.clear();
|
||
|
||
if (marks_count == 0)
|
||
{
|
||
/// Кусок пустой - все записи удалились.
|
||
Poco::File(part_path).remove(true);
|
||
checksums.files.clear();
|
||
return checksums;
|
||
}
|
||
|
||
{
|
||
/// Записываем файл с описанием столбцов.
|
||
WriteBufferFromFile out(part_path + "columns.txt", 4096);
|
||
columns_list.writeText(out);
|
||
}
|
||
|
||
{
|
||
/// Записываем файл с чексуммами.
|
||
WriteBufferFromFile out(part_path + "checksums.txt", 4096);
|
||
checksums.write(out);
|
||
}
|
||
|
||
return checksums;
|
||
}
|
||
|
||
MergeTreeData::DataPart::Index & getIndex()
|
||
{
|
||
return index_vec;
|
||
}
|
||
|
||
/// Сколько засечек уже записано.
|
||
size_t marksCount()
|
||
{
|
||
return marks_count;
|
||
}
|
||
|
||
private:
|
||
NamesAndTypesList columns_list;
|
||
String part_path;
|
||
|
||
size_t marks_count = 0;
|
||
|
||
SharedPtr<WriteBufferFromFile> index_file_stream;
|
||
SharedPtr<HashingWriteBuffer> index_stream;
|
||
MergeTreeData::DataPart::Index index_vec;
|
||
};
|
||
|
||
typedef Poco::SharedPtr<MergedBlockOutputStream> MergedBlockOutputStreamPtr;
|
||
|
||
/// Записывает только те, столбцы, что лежат в block
|
||
class MergedColumnOnlyOutputStream : public IMergedBlockOutputStream
|
||
{
|
||
public:
|
||
MergedColumnOnlyOutputStream(MergeTreeData & storage_, String part_path_, bool sync_, CompressionMethod compression_method)
|
||
: IMergedBlockOutputStream(
|
||
storage_, storage_.context.getSettings().min_compress_block_size,
|
||
storage_.context.getSettings().max_compress_block_size, compression_method),
|
||
part_path(part_path_), sync(sync_)
|
||
{
|
||
}
|
||
|
||
void write(const Block & block) override
|
||
{
|
||
if (!initialized)
|
||
{
|
||
column_streams.clear();
|
||
for (size_t i = 0; i < block.columns(); ++i)
|
||
{
|
||
addStream(part_path, block.getByPosition(i).name,
|
||
*block.getByPosition(i).type, 0, block.getByPosition(i).name);
|
||
}
|
||
initialized = true;
|
||
}
|
||
|
||
size_t rows = block.rows();
|
||
|
||
OffsetColumns offset_columns;
|
||
for (size_t i = 0; i < block.columns(); ++i)
|
||
{
|
||
const ColumnWithNameAndType & column = block.getByPosition(i);
|
||
writeData(column.name, *column.type, *column.column, offset_columns);
|
||
}
|
||
|
||
size_t written_for_last_mark = (storage.index_granularity - index_offset + rows) % storage.index_granularity;
|
||
index_offset = (storage.index_granularity - written_for_last_mark) % storage.index_granularity;
|
||
}
|
||
|
||
void writeSuffix() override
|
||
{
|
||
throw Exception("Method writeSuffix is not supported by MergedColumnOnlyOutputStream", ErrorCodes::NOT_IMPLEMENTED);
|
||
}
|
||
|
||
MergeTreeData::DataPart::Checksums writeSuffixAndGetChecksums()
|
||
{
|
||
MergeTreeData::DataPart::Checksums checksums;
|
||
|
||
for (auto & column_stream : column_streams)
|
||
{
|
||
column_stream.second->finalize();
|
||
if (sync)
|
||
column_stream.second->sync();
|
||
std::string column = escapeForFileName(column_stream.first);
|
||
column_stream.second->addToChecksums(checksums, column);
|
||
}
|
||
|
||
column_streams.clear();
|
||
initialized = false;
|
||
|
||
return checksums;
|
||
}
|
||
|
||
private:
|
||
String part_path;
|
||
|
||
bool initialized = false;
|
||
bool sync;
|
||
};
|
||
|
||
}
|