Fix mutations memory consumption finally

dbms/src/Interpreters/MutationsInterpreter.cpp

@@ -620,38 +620,23 @@ ASTPtr MutationsInterpreter::prepareInterpreterSelectQuery(std::vector<Stage> &
         }
         select->setExpression(ASTSelectQuery::Expression::WHERE, std::move(where_expression));
     }
-    auto metadata = storage->getInMemoryMetadata();
     /// We have to execute select in order of primary key
     /// because we don't sort results additionaly and don't have
     /// any guarantees on data order without ORDER BY. It's almost free, because we
     /// have optimization for data read in primary key order.
-    if (metadata.order_by_ast)
+    if (ASTPtr key_expr = storage->getSortingKeyAST(); key_expr && !key_expr->children.empty())
     {
         ASTPtr dummy;
-
+        auto res = std::make_shared<ASTExpressionList>();
-        ASTPtr key_expr;
+        for (const auto & key_part : key_expr->children)
-        if (metadata.primary_key_ast)
-            key_expr = metadata.primary_key_ast;
-        else
-            key_expr = metadata.order_by_ast;
-
-        bool empty = false;
-        /// In all other cases we cannot have empty key
-        if (auto key_function = key_expr->as<ASTFunction>())
-            empty = key_function->arguments->children.empty();
-
-        /// Not explicitely spicified empty key
-        if (!empty)
         {
             auto order_by_expr = std::make_shared<ASTOrderByElement>(1, 1, false, dummy, false, dummy, dummy, dummy);
+            order_by_expr->children.push_back(key_part);
 
-
-            order_by_expr->children.push_back(key_expr);
-            auto res = std::make_shared<ASTExpressionList>();
             res->children.push_back(order_by_expr);
-
-            select->setExpression(ASTSelectQuery::Expression::ORDER_BY, std::move(res));
         }
+
+        select->setExpression(ASTSelectQuery::Expression::ORDER_BY, std::move(res));
     }
 
     return select;

dbms/tests/queries/0_stateless/01200_mutations_memory_consumption.reference

@@ -1 +1,4 @@
 1
+1
+1
+1

dbms/tests/queries/0_stateless/01200_mutations_memory_consumption.sql

@@ -1,6 +1,6 @@
-DROP TABLE IF EXISTS table_with_pk;
+DROP TABLE IF EXISTS table_with_single_pk;
 
-CREATE TABLE table_with_pk
+CREATE TABLE table_with_single_pk
 (
   key UInt8,
   value String
@@ -8,9 +8,9 @@ CREATE TABLE table_with_pk
 ENGINE = MergeTree
 ORDER BY key;
 
-INSERT INTO table_with_pk SELECT number, toString(number % 10) FROM numbers(10000000);
+INSERT INTO table_with_single_pk SELECT number, toString(number % 10) FROM numbers(10000000);
 
-ALTER TABLE table_with_pk DELETE WHERE key % 77 = 0 SETTINGS mutations_sync = 1;
+ALTER TABLE table_with_single_pk DELETE WHERE key % 77 = 0 SETTINGS mutations_sync = 1;
 
 SYSTEM FLUSH LOGS;
 
@@ -20,6 +20,92 @@ SELECT
   DISTINCT read_bytes >= peak_memory_usage
 FROM
     system.part_log
-WHERE event_type = 'MutatePart' AND table = 'table_with_pk' AND database = currentDatabase();
+WHERE event_type = 'MutatePart' AND table = 'table_with_single_pk' AND database = currentDatabase();
 
-DROP TABLE IF EXISTS table_with_pk;
+DROP TABLE IF EXISTS table_with_single_pk;
+
+DROP TABLE IF EXISTS table_with_multi_pk;
+
+CREATE TABLE table_with_multi_pk
+(
+  key1 UInt8,
+  key2 UInt32,
+  key3 DateTime64(6, 'UTC'),
+  value String
+)
+ENGINE = MergeTree
+ORDER BY (key1, key2, key3);
+
+INSERT INTO table_with_multi_pk SELECT number % 32, number, toDateTime('2019-10-01 00:00:00'), toString(number % 10) FROM numbers(10000000);
+
+ALTER TABLE table_with_multi_pk DELETE WHERE key1 % 77 = 0 SETTINGS mutations_sync = 1;
+
+SYSTEM FLUSH LOGS;
+
+-- Memory usage for all mutations must be almost constant and less than
+-- read_bytes.
+SELECT
+  DISTINCT read_bytes >= peak_memory_usage
+  FROM
+      system.part_log
+ WHERE event_type = 'MutatePart' AND table = 'table_with_multi_pk' AND database = currentDatabase();
+
+DROP TABLE IF EXISTS table_with_multi_pk;
+
+
+DROP TABLE IF EXISTS table_with_function_pk;
+
+
+CREATE TABLE table_with_function_pk
+  (
+    key1 UInt8,
+    key2 UInt32,
+    key3 DateTime64(6, 'UTC'),
+    value String
+  )
+ENGINE = MergeTree
+ORDER BY (key1 + key2 + cast(value as UInt64), key2);
+
+INSERT INTO table_with_function_pk SELECT number % 32, number, toDateTime('2019-10-01 00:00:00'), toString(number % 10) FROM numbers(10000000);
+
+ALTER TABLE table_with_function_pk DELETE WHERE key1 % 77 = 0 SETTINGS mutations_sync = 1;
+
+SYSTEM FLUSH LOGS;
+
+-- Memory usage for all mutations must be almost constant and less than
+-- read_bytes.
+SELECT
+  DISTINCT read_bytes >= peak_memory_usage
+  FROM
+      system.part_log
+ WHERE event_type = 'MutatePart' AND table = 'table_with_function_pk' AND database = currentDatabase();
+
+DROP TABLE IF EXISTS table_with_function_pk;
+
+DROP TABLE IF EXISTS table_without_pk;
+
+CREATE TABLE table_without_pk
+(
+  key1 UInt8,
+  key2 UInt32,
+  key3 DateTime64(6, 'UTC'),
+  value String
+)
+ENGINE = MergeTree
+ORDER BY tuple();
+
+INSERT INTO table_without_pk SELECT number % 32, number, toDateTime('2019-10-01 00:00:00'), toString(number % 10) FROM numbers(10000000);
+
+ALTER TABLE table_without_pk DELETE WHERE key1 % 77 = 0 SETTINGS mutations_sync = 1;
+
+SYSTEM FLUSH LOGS;
+
+-- Memory usage for all mutations must be almost constant and less than
+-- read_bytes.
+SELECT
+  DISTINCT read_bytes >= peak_memory_usage
+  FROM
+      system.part_log
+ WHERE event_type = 'MutatePart' AND table = 'table_without_pk' AND database = currentDatabase();
+
+DROP TABLE IF EXISTS table_without_pk;