From 5d7a77c20771a803f34855f0e31fe58ad6566677 Mon Sep 17 00:00:00 2001
From: feng lv <fenglv15@mails.ucas.ac.cn>
Date: Tue, 22 Sep 2020 17:23:46 +0800
Subject: [PATCH] add mutation support for StorageMemory

---
 src/Storages/StorageMemory.cpp | 50 +++++++++++++++++++++++++++++++++-
 src/Storages/StorageMemory.h   |  2 ++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/src/Storages/StorageMemory.cpp b/src/Storages/StorageMemory.cpp
index 25e232dc4ad..9686fbe3062 100644
--- a/src/Storages/StorageMemory.cpp
+++ b/src/Storages/StorageMemory.cpp
@@ -2,8 +2,9 @@
 
 #include <DataStreams/IBlockInputStream.h>
 
-#include <Storages/StorageMemory.h>
+#include <Interpreters/MutationsInterpreter.h>
 #include <Storages/StorageFactory.h>
+#include <Storages/StorageMemory.h>
 
 #include <IO/WriteHelpers.h>
 #include <Processors/Sources/SourceWithProgress.h>
@@ -203,6 +204,53 @@ void StorageMemory::drop()
     data.clear();
 }
 
+static inline void columnUpdate(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::mutate(const MutationCommands & commands, const Context & context)
+{
+    auto metadata_snapshot_ = getInMemoryMetadataPtr();
+    auto storage_id_ = getStorageID();
+    auto storage_ptr_ = DatabaseCatalog::instance().getTable(storage_id_, context);
+    auto interpreter = std::make_unique<MutationsInterpreter>(storage_ptr_, metadata_snapshot_, commands, context, true);
+    auto in = interpreter->execute();
+
+    in->readPrefix();
+    BlocksList out;
+    Block block;
+    while (block = in->read())
+    {
+        out.push_back(block);
+    }
+    in->readSuffix();
+
+	std::lock_guard lock(mutex);
+
+    // all column affected
+    if (interpreter->isAffectingAllColumns())
+    {
+        std::swap(data, out);
+    }
+    else
+    {
+        auto data_it = data.begin();
+        auto out_it = out.begin();
+        while (data_it != data.end() && out_it != out.end())
+        {
+            columnUpdate(*data_it, *out_it);
+            ++data_it;
+            ++out_it;
+        }
+    }
+}
+
 void StorageMemory::truncate(
     const ASTPtr &, const StorageMetadataPtr &, const Context &, TableExclusiveLockHolder &)
 {
diff --git a/src/Storages/StorageMemory.h b/src/Storages/StorageMemory.h
index e67e3015028..f0d5d7aadec 100644
--- a/src/Storages/StorageMemory.h
+++ b/src/Storages/StorageMemory.h
@@ -43,6 +43,8 @@ public:
 
     void drop() override;
 
+    void mutate(const MutationCommands & commands, const Context & context) override;
+
     void truncate(const ASTPtr &, const StorageMetadataPtr &, const Context &, TableExclusiveLockHolder &) override;
 
     std::optional<UInt64> totalRows() const override;