#include #include #include #include #include #include #include #include #include #include #include #include namespace DB { namespace ErrorCodes { extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; } class MemorySource : public SourceWithProgress { using InitializerFunc = std::function &)>; public: MemorySource( Names column_names_, const StorageMemory & storage, const StorageMetadataPtr & metadata_snapshot, std::shared_ptr data_, std::shared_ptr> parallel_execution_index_, InitializerFunc initializer_func_ = {}) : SourceWithProgress(metadata_snapshot->getSampleBlockForColumns(column_names_, storage.getVirtuals(), storage.getStorageID())) , column_names_and_types(metadata_snapshot->getColumns().getByNames(ColumnsDescription::All, column_names_, true)) , data(data_) , parallel_execution_index(parallel_execution_index_) , initializer_func(std::move(initializer_func_)) { } String getName() const override { return "Memory"; } protected: Chunk generate() override { if (initializer_func) { initializer_func(data); initializer_func = {}; } size_t current_index = getAndIncrementExecutionIndex(); if (!data || current_index >= data->size()) { return {}; } const Block & src = (*data)[current_index]; Columns columns; columns.reserve(column_names_and_types.size()); /// Add only required columns to `res`. for (const auto & elem : column_names_and_types) columns.emplace_back(getColumnFromBlock(src, elem)); return Chunk(std::move(columns), src.rows()); } private: size_t getAndIncrementExecutionIndex() { if (parallel_execution_index) { return (*parallel_execution_index)++; } else { return execution_index++; } } const NamesAndTypesList column_names_and_types; size_t execution_index = 0; std::shared_ptr data; std::shared_ptr> parallel_execution_index; InitializerFunc initializer_func; }; class MemorySink : public SinkToStorage { public: MemorySink( StorageMemory & storage_, const StorageMetadataPtr & metadata_snapshot_) : SinkToStorage(metadata_snapshot_->getSampleBlock()) , storage(storage_) , metadata_snapshot(metadata_snapshot_) { } String getName() const override { return "MemorySink"; } void consume(Chunk chunk) override { auto block = getHeader().cloneWithColumns(chunk.getColumns()); metadata_snapshot->check(block, true); if (storage.compress) { Block compressed_block; for (const auto & elem : block) compressed_block.insert({ elem.column->compress(), elem.type, elem.name }); new_blocks.emplace_back(compressed_block); } else { new_blocks.emplace_back(block); } } void onFinish() override { size_t inserted_bytes = 0; size_t inserted_rows = 0; for (const auto & block : new_blocks) { inserted_bytes += block.allocatedBytes(); inserted_rows += block.rows(); } std::lock_guard lock(storage.mutex); auto new_data = std::make_unique(*(storage.data.get())); new_data->insert(new_data->end(), new_blocks.begin(), new_blocks.end()); storage.data.set(std::move(new_data)); storage.total_size_bytes.fetch_add(inserted_bytes, std::memory_order_relaxed); storage.total_size_rows.fetch_add(inserted_rows, std::memory_order_relaxed); } private: Blocks new_blocks; StorageMemory & storage; StorageMetadataPtr metadata_snapshot; }; StorageMemory::StorageMemory( const StorageID & table_id_, ColumnsDescription columns_description_, ConstraintsDescription constraints_, const String & comment, bool compress_) : IStorage(table_id_), data(std::make_unique()), compress(compress_) { StorageInMemoryMetadata storage_metadata; storage_metadata.setColumns(std::move(columns_description_)); storage_metadata.setConstraints(std::move(constraints_)); storage_metadata.setComment(comment); setInMemoryMetadata(storage_metadata); } Pipe StorageMemory::read( const Names & column_names, const StorageMetadataPtr & metadata_snapshot, SelectQueryInfo & /*query_info*/, ContextPtr /*context*/, QueryProcessingStage::Enum /*processed_stage*/, size_t /*max_block_size*/, unsigned num_streams) { metadata_snapshot->check(column_names, getVirtuals(), getStorageID()); if (delay_read_for_global_subqueries) { /// Note: for global subquery we use single source. /// Mainly, the reason is that at this point table is empty, /// and we don't know the number of blocks are going to be inserted into it. /// /// It may seem to be not optimal, but actually data from such table is used to fill /// set for IN or hash table for JOIN, which can't be done concurrently. /// Since no other manipulation with data is done, multiple sources shouldn't give any profit. return Pipe(std::make_shared( column_names, *this, metadata_snapshot, nullptr /* data */, nullptr /* parallel execution index */, [this](std::shared_ptr & data_to_initialize) { data_to_initialize = data.get(); })); } auto current_data = data.get(); size_t size = current_data->size(); if (num_streams > size) num_streams = size; Pipes pipes; auto parallel_execution_index = std::make_shared>(0); for (size_t stream = 0; stream < num_streams; ++stream) { pipes.emplace_back(std::make_shared(column_names, *this, metadata_snapshot, current_data, parallel_execution_index)); } return Pipe::unitePipes(std::move(pipes)); } SinkToStoragePtr StorageMemory::write(const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, ContextPtr /*context*/) { return std::make_shared(*this, metadata_snapshot); } void StorageMemory::drop() { data.set(std::make_unique()); total_size_bytes.store(0, std::memory_order_relaxed); total_size_rows.store(0, std::memory_order_relaxed); } static inline void updateBlockData(Block & old_block, const Block & new_block) { for (const auto & it : new_block) { auto col_name = it.name; auto & col_with_type_name = old_block.getByName(col_name); col_with_type_name.column = it.column; } } void StorageMemory::checkMutationIsPossible(const MutationCommands & /*commands*/, const Settings & /*settings*/) const { /// Some validation will be added } void StorageMemory::mutate(const MutationCommands & commands, ContextPtr context) { std::lock_guard lock(mutex); auto metadata_snapshot = getInMemoryMetadataPtr(); auto storage = getStorageID(); auto storage_ptr = DatabaseCatalog::instance().getTable(storage, context); /// When max_threads > 1, the order of returning blocks is uncertain, /// which will lead to inconsistency after updateBlockData. auto new_context = Context::createCopy(context); new_context->setSetting("max_streams_to_max_threads_ratio", 1); new_context->setSetting("max_threads", 1); auto interpreter = std::make_unique(storage_ptr, metadata_snapshot, commands, new_context, true); auto pipeline = interpreter->execute(); PullingPipelineExecutor executor(pipeline); Blocks out; Block block; while (executor.pull(block)) { if (compress) for (auto & elem : block) elem.column = elem.column->compress(); out.push_back(block); } std::unique_ptr new_data; // all column affected if (interpreter->isAffectingAllColumns()) { new_data = std::make_unique(out); } else { /// just some of the column affected, we need update it with new column new_data = std::make_unique(*(data.get())); auto data_it = new_data->begin(); auto out_it = out.begin(); while (data_it != new_data->end()) { /// Mutation does not change the number of blocks assert(out_it != out.end()); updateBlockData(*data_it, *out_it); ++data_it; ++out_it; } assert(out_it == out.end()); } size_t rows = 0; size_t bytes = 0; for (const auto & buffer : *new_data) { rows += buffer.rows(); bytes += buffer.bytes(); } total_size_bytes.store(bytes, std::memory_order_relaxed); total_size_rows.store(rows, std::memory_order_relaxed); data.set(std::move(new_data)); } void StorageMemory::truncate( const ASTPtr &, const StorageMetadataPtr &, ContextPtr, TableExclusiveLockHolder &) { data.set(std::make_unique()); total_size_bytes.store(0, std::memory_order_relaxed); total_size_rows.store(0, std::memory_order_relaxed); } std::optional StorageMemory::totalRows(const Settings &) const { /// All modifications of these counters are done under mutex which automatically guarantees synchronization/consistency /// When run concurrently we are fine with any value: "before" or "after" return total_size_rows.load(std::memory_order_relaxed); } std::optional StorageMemory::totalBytes(const Settings &) const { return total_size_bytes.load(std::memory_order_relaxed); } void registerStorageMemory(StorageFactory & factory) { factory.registerStorage("Memory", [](const StorageFactory::Arguments & args) { if (!args.engine_args.empty()) throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Engine {} doesn't support any arguments ({} given)", args.engine_name, args.engine_args.size()); bool has_settings = args.storage_def->settings; MemorySettings settings; if (has_settings) settings.loadFromQuery(*args.storage_def); return StorageMemory::create(args.table_id, args.columns, args.constraints, args.comment, settings.compress); }, { .supports_settings = true, .supports_parallel_insert = true, }); } }