From ef81f3313f01389da3fed91563a6d42084741281 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Mar=C3=ADn?= <git@rmr.ninja>
Date: Mon, 6 May 2024 20:58:04 +0200
Subject: [PATCH] Optimize ActionsDAG::updateHeader

---
 src/Interpreters/ActionsDAG.cpp | 32 ++++++++++++++++++--------------
 src/Interpreters/ActionsDAG.h   |  2 +-
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/src/Interpreters/ActionsDAG.cpp b/src/Interpreters/ActionsDAG.cpp
index 30eb908330b..91e75dca2e1 100644
--- a/src/Interpreters/ActionsDAG.cpp
+++ b/src/Interpreters/ActionsDAG.cpp
@@ -708,27 +708,28 @@ static ColumnWithTypeAndName executeActionForPartialResult(const ActionsDAG::Nod
     return res_column;
 }
 
-Block ActionsDAG::updateHeader(Block header) const
+Block ActionsDAG::updateHeader(const Block & header) const
 {
     IntermediateExecutionResult node_to_column;
     std::set<size_t> pos_to_remove;
 
     {
-        std::unordered_map<std::string_view, std::list<size_t>> input_positions;
+        size_t out = inputs.size(); /// Always out of range
+        std::unordered_map<std::string_view, size_t> input_positions;
+        input_positions.reserve(inputs.size());
 
         for (size_t pos = 0; pos < inputs.size(); ++pos)
-            input_positions[inputs[pos]->result_name].emplace_back(pos);
+            input_positions[inputs[pos]->result_name] = pos;
 
         for (size_t pos = 0; pos < header.columns(); ++pos)
         {
             const auto & col = header.getByPosition(pos);
             auto it = input_positions.find(col.name);
-            if (it != input_positions.end() && !it->second.empty())
+            if (it != input_positions.end() && it->second != out)
             {
-                auto & list = it->second;
                 pos_to_remove.insert(pos);
-                node_to_column[inputs[list.front()]] = col;
-                list.pop_front();
+                node_to_column[inputs[it->second]] = col;
+                it->second = out;
             }
         }
     }
@@ -746,18 +747,21 @@ Block ActionsDAG::updateHeader(Block header) const
         throw;
     }
 
-    if (isInputProjected())
-        header.clear();
-    else
-        header.erase(pos_to_remove);
 
     Block res;
-
+    res.reserve(result_columns.size());
     for (auto & col : result_columns)
         res.insert(std::move(col));
 
-    for (auto && item : header)
-        res.insert(std::move(item));
+    if (isInputProjected())
+        return res;
+
+    res.reserve(header.columns() - pos_to_remove.size());
+    for (size_t i = 0; i < header.columns(); i++)
+    {
+        if (!pos_to_remove.contains(i))
+            res.insert(header.data[i]);
+    }
 
     return res;
 }
diff --git a/src/Interpreters/ActionsDAG.h b/src/Interpreters/ActionsDAG.h
index 8c0bcf8fdc0..64e6a0a7998 100644
--- a/src/Interpreters/ActionsDAG.h
+++ b/src/Interpreters/ActionsDAG.h
@@ -272,7 +272,7 @@ public:
     ///
     /// In addition, check that result constants are constants according to DAG.
     /// In case if function return constant, but arguments are not constant, materialize it.
-    Block updateHeader(Block header) const;
+    Block updateHeader(const Block & header) const;
 
     using IntermediateExecutionResult = std::unordered_map<const Node *, ColumnWithTypeAndName>;
     static ColumnsWithTypeAndName evaluatePartialResult(