From 18491d89206c8c2d95c35c6ff4b3292bd769238f Mon Sep 17 00:00:00 2001
From: feng lv <fenglv15@mails.ucas.ac.cn>
Date: Mon, 23 Nov 2020 05:30:36 +0000
Subject: [PATCH] fix

---
 src/Storages/StorageMemory.cpp                | 41 ++++++++-----------
 ...1498_alter_column_storage_memory.reference |  0
 .../01498_alter_column_storage_memory.sql     | 11 +++++
 3 files changed, 29 insertions(+), 23 deletions(-)
 create mode 100644 tests/queries/0_stateless/01498_alter_column_storage_memory.reference
 create mode 100644 tests/queries/0_stateless/01498_alter_column_storage_memory.sql

diff --git a/src/Storages/StorageMemory.cpp b/src/Storages/StorageMemory.cpp
index b6b6cb6285b..f9227764792 100644
--- a/src/Storages/StorageMemory.cpp
+++ b/src/Storages/StorageMemory.cpp
@@ -207,7 +207,6 @@ BlockOutputStreamPtr StorageMemory::write(const ASTPtr & /*query*/, const Storag
 
 void StorageMemory::drop()
 {
-    std::lock_guard lock(mutex);
     data.set(std::make_unique<BlocksList>());
     total_size_bytes.store(0, std::memory_order_relaxed);
     total_size_rows.store(0, std::memory_order_relaxed);
@@ -241,49 +240,45 @@ void StorageMemory::mutate(const MutationCommands & commands, const Context & co
     }
     in->readSuffix();
 
+    std::unique_ptr<BlocksList> new_data;
+
     // all column affected
     if (interpreter->isAffectingAllColumns())
     {
-        size_t rows = 0;
-        size_t bytes = 0;
-        for (const auto & buffer : out)
-        {
-            rows += buffer.rows();
-            bytes += buffer.bytes();
-        }
-        data.set(std::make_unique<BlocksList>(out));
-        total_size_bytes.store(rows, std::memory_order_relaxed);
-        total_size_rows.store(bytes, std::memory_order_relaxed);
+        new_data = std::make_unique<BlocksList>(out);
     }
     else
     {
-        auto new_data = std::make_unique<BlocksList>(*(data.get()));
+        /// just some of the column affected, we need update it with new column
+        new_data = std::make_unique<BlocksList>(*(data.get()));
         auto data_it = new_data->begin();
         auto out_it = out.begin();
+        /// Mutation does not change the number of blocks, so we don't need
+        // to check whether old data and new data have same number of blocks or not
         while (data_it != new_data->end() && out_it != out.end())
         {
             updateBlockData(*data_it, *out_it);
             ++data_it;
             ++out_it;
         }
-        size_t rows = 0;
-        size_t bytes = 0;
-        for (const auto & buffer : *new_data)
-        {
-            rows += buffer.rows();
-            bytes += buffer.bytes();
-        }
-        total_size_bytes.store(rows, std::memory_order_relaxed);
-        total_size_rows.store(bytes, std::memory_order_relaxed);
-        data.set(std::move(new_data));
     }
+
+    size_t rows = 0;
+    size_t bytes = 0;
+    for (const auto & buffer : *new_data)
+    {
+        rows += buffer.rows();
+        bytes += buffer.bytes();
+    }
+    total_size_bytes.store(rows, std::memory_order_relaxed);
+    total_size_rows.store(bytes, std::memory_order_relaxed);
+    data.set(std::move(new_data));
 }
 
 
 void StorageMemory::truncate(
     const ASTPtr &, const StorageMetadataPtr &, const Context &, TableExclusiveLockHolder &)
 {
-    std::lock_guard lock(mutex);
     data.set(std::make_unique<BlocksList>());
     total_size_bytes.store(0, std::memory_order_relaxed);
     total_size_rows.store(0, std::memory_order_relaxed);
diff --git a/tests/queries/0_stateless/01498_alter_column_storage_memory.reference b/tests/queries/0_stateless/01498_alter_column_storage_memory.reference
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01498_alter_column_storage_memory.sql b/tests/queries/0_stateless/01498_alter_column_storage_memory.sql
new file mode 100644
index 00000000000..5f0b001f178
--- /dev/null
+++ b/tests/queries/0_stateless/01498_alter_column_storage_memory.sql
@@ -0,0 +1,11 @@
+DROP TABLE IF EXISTS defaults;
+CREATE TABLE defaults
+(
+	n Int32,
+	s String
+)ENGINE = Memory();
+
+ALTER TABLE defaults ADD COLUMN m Int8; -- { serverError 48 }
+ALTER TABLE defaults DROP COLUMN n; -- { serverError 48 }
+
+DROP TABLE defaults;