void test1()
{
Context context;
StoragePtr table = StorageSystemNumbers::create("numbers");
Names column_names;
column_names.push_back("number");
QueryProcessingStage::Enum stage1;
QueryProcessingStage::Enum stage2;
QueryProcessingStage::Enum stage3;
BlockInputStreams streams;
streams.emplace_back(std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage1, 1)[0], 30, 30000));
streams.emplace_back(std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage2, 1)[0], 30, 2000));
streams.emplace_back(std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage3, 1)[0], 30, 100));
UnionBlockInputStream<> union_stream(streams, nullptr, 2);
WriteBufferFromFileDescriptor wb(STDERR_FILENO);
Block sample = table->getSampleBlock();
BlockOutputStreamPtr out = context.getOutputFormat("TabSeparated", wb, sample);
while (Block block = union_stream.read())
{
out->write(block);
wb.next();
}
//copyData(union_stream, *out);
}
void inputThread(BlockInputStreamPtr in, BlockOutputStreamPtr out, WriteBuffer & wb, std::mutex & mutex)
{
while (Block block = in->read())
{
std::lock_guard<std::mutex> lock(mutex);
out->write(block);
wb.next();
}
}
int main(int argc, char ** argv)
try
{
using namespace DB;
const size_t rows = 10000000;
/// создаём таблицу с парой столбцов
NamesAndTypesListPtr names_and_types = std::make_shared<NamesAndTypesList>();
names_and_types->push_back(NameAndTypePair("a", std::make_shared<DataTypeUInt64>()));
names_and_types->push_back(NameAndTypePair("b", std::make_shared<DataTypeUInt8>()));
StoragePtr table = StorageLog::create("./", "test", names_and_types);
/// пишем в неё
{
Block block;
ColumnWithTypeAndName column1;
column1.name = "a";
column1.type = table->getDataTypeByName("a");
column1.column = column1.type->createColumn();
ColumnUInt64::Container_t & vec1 = typeid_cast<ColumnUInt64&>(*column1.column).getData();
vec1.resize(rows);
for (size_t i = 0; i < rows; ++i)
vec1[i] = i;
block.insert(column1);
ColumnWithTypeAndName column2;
column2.name = "b";
column2.type = table->getDataTypeByName("b");
column2.column = column2.type->createColumn();
ColumnUInt8::Container_t & vec2 = typeid_cast<ColumnUInt8&>(*column2.column).getData();
vec2.resize(rows);
for (size_t i = 0; i < rows; ++i)
vec2[i] = i * 2;
block.insert(column2);
BlockOutputStreamPtr out = table->write({}, {});
out->write(block);
}
/// читаем из неё
{
Names column_names;
column_names.push_back("a");
column_names.push_back("b");
QueryProcessingStage::Enum stage;
BlockInputStreamPtr in = table->read(column_names, 0, Context{}, Settings(), stage)[0];
Block sample;
{
ColumnWithTypeAndName col;
col.type = std::make_shared<DataTypeUInt64>();
sample.insert(col);
}
{
ColumnWithTypeAndName col;
col.type = std::make_shared<DataTypeUInt8>();
sample.insert(col);
}
WriteBufferFromOStream out_buf(std::cout);
LimitBlockInputStream in_limit(in, 10, 0);
RowOutputStreamPtr output_ = std::make_shared<TabSeparatedRowOutputStream>(out_buf, sample);
BlockOutputStreamFromRowOutputStream output(output_);
copyData(in_limit, output);
}
return 0;
}
catch (const DB::Exception & e)
{
std::cerr << e.what() << ", " << e.displayText() << std::endl;
return 1;
}
void executeQuery(
ReadBuffer & istr,
WriteBuffer & ostr,
Context & context,
BlockInputStreamPtr & query_plan,
std::function<void(const String &)> set_content_type)
{
PODArray<char> parse_buf;
const char * begin;
const char * end;
/// If 'istr' is empty now, fetch next data into buffer.
if (istr.buffer().size() == 0)
istr.next();
size_t max_query_size = context.getSettingsRef().max_query_size;
if (istr.buffer().end() - istr.position() >= static_cast<ssize_t>(max_query_size))
{
/// If remaining buffer space in 'istr' is enough to parse query up to 'max_query_size' bytes, then parse inplace.
begin = istr.position();
end = istr.buffer().end();
istr.position() += end - begin;
}
else
{
/// If not - copy enough data into 'parse_buf'.
parse_buf.resize(max_query_size);
parse_buf.resize(istr.read(&parse_buf[0], max_query_size));
begin = &parse_buf[0];
end = begin + parse_buf.size();
}
ASTPtr ast;
BlockIO streams;
std::tie(ast, streams) = executeQueryImpl(begin, end, context, false, QueryProcessingStage::Complete);
try
{
if (streams.out)
{
const ASTInsertQuery * ast_insert_query = dynamic_cast<const ASTInsertQuery *>(ast.get());
if (!ast_insert_query)
throw Exception("Logical error: query requires data to insert, but it is not INSERT query", ErrorCodes::LOGICAL_ERROR);
String format = ast_insert_query->format;
if (format.empty())
format = "Values";
/// Data could be in parsed (ast_insert_query.data) and in not parsed yet (istr) part of query.
ConcatReadBuffer::ReadBuffers buffers;
ReadBuffer buf1(const_cast<char *>(ast_insert_query->data), ast_insert_query->data ? ast_insert_query->end - ast_insert_query->data : 0, 0);
if (ast_insert_query->data)
buffers.push_back(&buf1);
buffers.push_back(&istr);
/** NOTE Must not read from 'istr' before read all between 'ast_insert_query.data' and 'ast_insert_query.end'.
* - because 'query.data' could refer to memory piece, used as buffer for 'istr'.
*/
ConcatReadBuffer data_istr(buffers);
BlockInputStreamPtr in{
context.getInputFormat(
format, data_istr, streams.out_sample, context.getSettings().max_insert_block_size)};
copyData(*in, *streams.out);
}
if (streams.in)
{
const ASTQueryWithOutput * ast_query_with_output = dynamic_cast<const ASTQueryWithOutput *>(ast.get());
String format_name = ast_query_with_output && (ast_query_with_output->getFormat() != nullptr)
? typeid_cast<const ASTIdentifier &>(*ast_query_with_output->getFormat()).name
: context.getDefaultFormat();
BlockOutputStreamPtr out = context.getOutputFormat(format_name, ostr, streams.in_sample);
if (IProfilingBlockInputStream * stream = dynamic_cast<IProfilingBlockInputStream *>(streams.in.get()))
{
/// NOTE Progress callback takes shared ownership of 'out'.
stream->setProgressCallback([out] (const Progress & progress) { out->onProgress(progress); });
}
if (set_content_type)
set_content_type(out->getContentType());
copyData(*streams.in, *out);
}
}
catch (...)
{
streams.onException();
throw;
}
streams.onFinish();
}
void executeQuery(
ReadBuffer & istr,
WriteBuffer & ostr,
bool allow_into_outfile,
Context & context,
std::function<void(const String &)> set_content_type)
{
PODArray<char> parse_buf;
const char * begin;
const char * end;
/// If 'istr' is empty now, fetch next data into buffer.
if (istr.buffer().size() == 0)
istr.next();
size_t max_query_size = context.getSettingsRef().max_query_size;
if (istr.buffer().end() - istr.position() >= static_cast<ssize_t>(max_query_size))
{
/// If remaining buffer space in 'istr' is enough to parse query up to 'max_query_size' bytes, then parse inplace.
begin = istr.position();
end = istr.buffer().end();
istr.position() += end - begin;
}
else
{
/// If not - copy enough data into 'parse_buf'.
parse_buf.resize(max_query_size);
parse_buf.resize(istr.read(&parse_buf[0], max_query_size));
begin = &parse_buf[0];
end = begin + parse_buf.size();
}
ASTPtr ast;
BlockIO streams;
std::tie(ast, streams) = executeQueryImpl(begin, end, context, false, QueryProcessingStage::Complete);
try
{
if (streams.out)
{
InputStreamFromASTInsertQuery in(ast, istr, streams, context);
copyData(in, *streams.out);
}
if (streams.in)
{
const ASTQueryWithOutput * ast_query_with_output = dynamic_cast<const ASTQueryWithOutput *>(ast.get());
WriteBuffer * out_buf = &ostr;
std::experimental::optional<WriteBufferFromFile> out_file_buf;
if (ast_query_with_output && ast_query_with_output->out_file)
{
if (!allow_into_outfile)
throw Exception("INTO OUTFILE is not allowed", ErrorCodes::INTO_OUTFILE_NOT_ALLOWED);
const auto & out_file = typeid_cast<const ASTLiteral &>(*ast_query_with_output->out_file).value.safeGet<std::string>();
out_file_buf.emplace(out_file, DBMS_DEFAULT_BUFFER_SIZE, O_WRONLY | O_EXCL | O_CREAT);
out_buf = &out_file_buf.value();
}
String format_name = ast_query_with_output && (ast_query_with_output->format != nullptr)
? typeid_cast<const ASTIdentifier &>(*ast_query_with_output->format).name
: context.getDefaultFormat();
BlockOutputStreamPtr out = context.getOutputFormat(format_name, *out_buf, streams.in_sample);
if (auto stream = dynamic_cast<IProfilingBlockInputStream *>(streams.in.get()))
{
/// Save previous progress callback if any. TODO Do it more conveniently.
auto previous_progress_callback = context.getProgressCallback();
/// NOTE Progress callback takes shared ownership of 'out'.
stream->setProgressCallback([out, previous_progress_callback] (const Progress & progress)
{
if (previous_progress_callback)
previous_progress_callback(progress);
out->onProgress(progress);
});
}
if (set_content_type)
set_content_type(out->getContentType());
copyData(*streams.in, *out);
}
}
catch (...)
{
streams.onException();
throw;
}
streams.onFinish();
}
void test2()
{
Context context;
StoragePtr table = StorageSystemNumbers::create("numbers");
Names column_names;
column_names.push_back("number");
QueryProcessingStage::Enum stage1;
QueryProcessingStage::Enum stage2;
QueryProcessingStage::Enum stage3;
BlockExtraInfo extra_info1;
extra_info1.host = "host1";
extra_info1.resolved_address = "127.0.0.1";
extra_info1.port = 9000;
extra_info1.user = "******";
BlockExtraInfo extra_info2;
extra_info2.host = "host2";
extra_info2.resolved_address = "127.0.0.2";
extra_info2.port = 9001;
extra_info2.user = "******";
BlockExtraInfo extra_info3;
extra_info3.host = "host3";
extra_info3.resolved_address = "127.0.0.3";
extra_info3.port = 9003;
extra_info3.user = "******";
BlockInputStreams streams;
BlockInputStreamPtr stream1 = std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage1, 1)[0], 30, 30000);
stream1 = std::make_shared<BlockExtraInfoInputStream>(stream1, extra_info1);
streams.emplace_back(stream1);
BlockInputStreamPtr stream2 = std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage2, 1)[0], 30, 2000);
stream2 = std::make_shared<BlockExtraInfoInputStream>(stream2, extra_info2);
streams.emplace_back(stream2);
BlockInputStreamPtr stream3 = std::make_shared<LimitBlockInputStream>(table->read(column_names, 0, context, Settings(), stage3, 1)[0], 30, 100);
stream3 = std::make_shared<BlockExtraInfoInputStream>(stream3, extra_info3);
streams.emplace_back(stream3);
UnionBlockInputStream<StreamUnionMode::ExtraInfo> union_stream(streams, nullptr, 2);
auto getSampleBlock = []()
{
Block block;
ColumnWithTypeAndName col;
col.name = "number";
col.type = std::make_shared<DataTypeUInt64>();
col.column = col.type->createColumn();
block.insert(col);
col.name = "host_name";
col.type = std::make_shared<DataTypeString>();
col.column = col.type->createColumn();
block.insert(col);
col.name = "host_address";
col.type = std::make_shared<DataTypeString>();
col.column = col.type->createColumn();
block.insert(col);
col.name = "port";
col.type = std::make_shared<DataTypeUInt16>();
col.column = col.type->createColumn();
block.insert(col);
col.name = "user";
col.type = std::make_shared<DataTypeString>();
col.column = col.type->createColumn();
block.insert(col);
return block;
};
WriteBufferFromFileDescriptor wb(STDERR_FILENO);
Block sample = getSampleBlock();
BlockOutputStreamPtr out = context.getOutputFormat("TabSeparated", wb, sample);
while (Block block = union_stream.read())
{
const auto & col = block.safeGetByPosition(0);
auto extra_info = union_stream.getBlockExtraInfo();
ColumnPtr host_name_column = std::make_shared<ColumnString>();
ColumnPtr host_address_column = std::make_shared<ColumnString>();
ColumnPtr port_column = std::make_shared<ColumnUInt16>();
ColumnPtr user_column = std::make_shared<ColumnString>();
size_t row_count = block.rows();
for (size_t i = 0; i < row_count; ++i)
{
host_name_column->insert(extra_info.resolved_address);
host_address_column->insert(extra_info.host);
port_column->insert(static_cast<UInt64>(extra_info.port));
user_column->insert(extra_info.user);
}
Block out_block;
out_block.insert(ColumnWithTypeAndName(col.column->clone(), col.type, col.name));
out_block.insert(ColumnWithTypeAndName(host_name_column, std::make_shared<DataTypeString>(), "host_name"));
out_block.insert(ColumnWithTypeAndName(host_address_column, std::make_shared<DataTypeString>(), "host_address"));
out_block.insert(ColumnWithTypeAndName(port_column, std::make_shared<DataTypeUInt16>(), "port"));
out_block.insert(ColumnWithTypeAndName(user_column, std::make_shared<DataTypeString>(), "user"));
out->write(out_block);
wb.next();
}
//copyData(union_stream, *out);
}
