mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-13 09:52:38 +00:00
191 lines
4.8 KiB
C++
191 lines
4.8 KiB
C++
#include <Common/Exception.h>
|
|
|
|
#include <DataStreams/IBlockInputStream.h>
|
|
|
|
#include <Storages/StorageMemory.h>
|
|
#include <Storages/StorageFactory.h>
|
|
|
|
#include <IO/WriteHelpers.h>
|
|
#include <Processors/Sources/SourceWithProgress.h>
|
|
#include <Processors/Pipe.h>
|
|
|
|
|
|
namespace DB
|
|
{
|
|
|
|
namespace ErrorCodes
|
|
{
|
|
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
|
}
|
|
|
|
|
|
class MemorySource : public SourceWithProgress
|
|
{
|
|
public:
|
|
MemorySource(
|
|
Names column_names_,
|
|
BlocksList::iterator begin_,
|
|
BlocksList::iterator end_,
|
|
const StorageMemory & storage,
|
|
const StorageMetadataPtr & metadata_snapshot)
|
|
: SourceWithProgress(metadata_snapshot->getSampleBlockForColumns(column_names_, storage.getVirtuals(), storage.getStorageID()))
|
|
, column_names(std::move(column_names_))
|
|
, begin(begin_)
|
|
, end(end_)
|
|
, it(begin)
|
|
{
|
|
}
|
|
|
|
String getName() const override { return "Memory"; }
|
|
|
|
protected:
|
|
Chunk generate() override
|
|
{
|
|
if (it == end)
|
|
{
|
|
return {};
|
|
}
|
|
else
|
|
{
|
|
Block src = *it;
|
|
Columns columns;
|
|
columns.reserve(column_names.size());
|
|
|
|
/// Add only required columns to `res`.
|
|
for (const auto & name : column_names)
|
|
columns.emplace_back(src.getByName(name).column);
|
|
|
|
++it;
|
|
return Chunk(std::move(columns), src.rows());
|
|
}
|
|
}
|
|
private:
|
|
Names column_names;
|
|
BlocksList::iterator begin;
|
|
BlocksList::iterator end;
|
|
BlocksList::iterator it;
|
|
};
|
|
|
|
|
|
class MemoryBlockOutputStream : public IBlockOutputStream
|
|
{
|
|
public:
|
|
explicit MemoryBlockOutputStream(
|
|
StorageMemory & storage_,
|
|
const StorageMetadataPtr & metadata_snapshot_)
|
|
: storage(storage_)
|
|
, metadata_snapshot(metadata_snapshot_)
|
|
{}
|
|
|
|
Block getHeader() const override { return metadata_snapshot->getSampleBlock(); }
|
|
|
|
void write(const Block & block) override
|
|
{
|
|
metadata_snapshot->check(block, true);
|
|
std::lock_guard lock(storage.mutex);
|
|
storage.data.push_back(block);
|
|
}
|
|
private:
|
|
StorageMemory & storage;
|
|
StorageMetadataPtr metadata_snapshot;
|
|
};
|
|
|
|
|
|
StorageMemory::StorageMemory(const StorageID & table_id_, ColumnsDescription columns_description_, ConstraintsDescription constraints_)
|
|
: IStorage(table_id_)
|
|
{
|
|
StorageInMemoryMetadata storage_metadata;
|
|
storage_metadata.setColumns(std::move(columns_description_));
|
|
storage_metadata.setConstraints(std::move(constraints_));
|
|
setInMemoryMetadata(storage_metadata);
|
|
}
|
|
|
|
|
|
Pipes StorageMemory::read(
|
|
const Names & column_names,
|
|
const StorageMetadataPtr & metadata_snapshot,
|
|
const SelectQueryInfo & /*query_info*/,
|
|
const Context & /*context*/,
|
|
QueryProcessingStage::Enum /*processed_stage*/,
|
|
size_t /*max_block_size*/,
|
|
unsigned num_streams)
|
|
{
|
|
metadata_snapshot->check(column_names, getVirtuals(), getStorageID());
|
|
|
|
std::lock_guard lock(mutex);
|
|
|
|
size_t size = data.size();
|
|
|
|
if (num_streams > size)
|
|
num_streams = size;
|
|
|
|
Pipes pipes;
|
|
|
|
for (size_t stream = 0; stream < num_streams; ++stream)
|
|
{
|
|
BlocksList::iterator begin = data.begin();
|
|
BlocksList::iterator end = data.begin();
|
|
|
|
std::advance(begin, stream * size / num_streams);
|
|
std::advance(end, (stream + 1) * size / num_streams);
|
|
|
|
pipes.emplace_back(std::make_shared<MemorySource>(column_names, begin, end, *this, metadata_snapshot));
|
|
}
|
|
|
|
return pipes;
|
|
}
|
|
|
|
|
|
BlockOutputStreamPtr StorageMemory::write(const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, const Context & /*context*/)
|
|
{
|
|
return std::make_shared<MemoryBlockOutputStream>(*this, metadata_snapshot);
|
|
}
|
|
|
|
|
|
void StorageMemory::drop()
|
|
{
|
|
std::lock_guard lock(mutex);
|
|
data.clear();
|
|
}
|
|
|
|
void StorageMemory::truncate(
|
|
const ASTPtr &, const StorageMetadataPtr &, const Context &, TableExclusiveLockHolder &)
|
|
{
|
|
std::lock_guard lock(mutex);
|
|
data.clear();
|
|
}
|
|
|
|
std::optional<UInt64> StorageMemory::totalRows() const
|
|
{
|
|
UInt64 rows = 0;
|
|
std::lock_guard lock(mutex);
|
|
for (const auto & buffer : data)
|
|
rows += buffer.rows();
|
|
return rows;
|
|
}
|
|
|
|
std::optional<UInt64> StorageMemory::totalBytes() const
|
|
{
|
|
UInt64 bytes = 0;
|
|
std::lock_guard lock(mutex);
|
|
for (const auto & buffer : data)
|
|
bytes += buffer.allocatedBytes();
|
|
return bytes;
|
|
}
|
|
|
|
|
|
void registerStorageMemory(StorageFactory & factory)
|
|
{
|
|
factory.registerStorage("Memory", [](const StorageFactory::Arguments & args)
|
|
{
|
|
if (!args.engine_args.empty())
|
|
throw Exception(
|
|
"Engine " + args.engine_name + " doesn't support any arguments (" + toString(args.engine_args.size()) + " given)",
|
|
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
|
|
|
return StorageMemory::create(args.table_id, args.columns, args.constraints);
|
|
});
|
|
}
|
|
|
|
}
|