Update sort desc: more efficient original node search in ActionsDAG

src/Interpreters/ActionsDAG.cpp

@@ -2500,8 +2500,21 @@ ActionsDAGPtr ActionsDAG::buildFilterActionsDAG(
 FindOriginalNodeForOutputName::FindOriginalNodeForOutputName(const ActionsDAGPtr & actions_)
     :actions(actions_)
 {
-    for (const auto * node : actions->getOutputs())
-        index.emplace(node->result_name, node);
+    const auto & actions_outputs = actions->getOutputs();
+    for (const auto * output_node : actions_outputs)
+    {
+        /// find input node which refers to the output node
+        /// consider only aliases on the path
+        const auto * node = output_node;
+        while (node && node->type == ActionsDAG::ActionType::ALIAS)
+        {
+            /// alias has only one child
+            chassert(node->children.size() == 1);
+            node = node->children.front();
+        }
+        if (node && node->type == ActionsDAG::ActionType::INPUT)
+            index.emplace(output_node->result_name, node);
+    }
 }
 
 const ActionsDAG::Node * FindOriginalNodeForOutputName::find(const String & output_name)
@@ -2510,17 +2523,36 @@ const ActionsDAG::Node * FindOriginalNodeForOutputName::find(const String & outp
     if (it == index.end())
         return nullptr;
 
-    /// find original(non alias) node it refers to
-    const ActionsDAG::Node * node = it->second;
-    while (node && node->type == ActionsDAG::ActionType::ALIAS)
+    return it->second;
+}
+
+FindAliasForInputName::FindAliasForInputName(const ActionsDAGPtr & actions_)
+    :actions(actions_)
+{
+    const auto & actions_outputs = actions->getOutputs();
+    for (const auto * output_node : actions_outputs)
     {
-        chassert(!node->children.empty());
-        node = node->children.front();
+        /// find input node which corresponds to alias
+        const auto * node = output_node;
+        while (node && node->type == ActionsDAG::ActionType::ALIAS)
+        {
+            /// alias has only one child
+            chassert(node->children.size() == 1);
+            node = node->children.front();
+        }
+        if (node && node->type == ActionsDAG::ActionType::INPUT)
+            /// node can have several aliases but we consider only the first one
+            index.emplace(node->result_name, output_node);
     }
-    if (node && node->type != ActionsDAG::ActionType::INPUT)
+}
+
+const ActionsDAG::Node * FindAliasForInputName::find(const String & name)
+{
+    const auto it = index.find(name);
+    if (it == index.end())
         return nullptr;
 
-    return node;
+    return it->second;
 }
 
 }

src/Interpreters/ActionsDAG.h

@@ -408,7 +408,20 @@ class FindOriginalNodeForOutputName
 
 public:
     explicit FindOriginalNodeForOutputName(const ActionsDAGPtr & actions);
-    const ActionsDAG::Node* find(const String& output_name);
+    const ActionsDAG::Node * find(const String & output_name);
+
+private:
+    ActionsDAGPtr actions;
+    NameToNodeIndex index;
+};
+
+class FindAliasForInputName
+{
+    using NameToNodeIndex = std::unordered_map<std::string_view, const ActionsDAG::Node *>;
+
+public:
+    explicit FindAliasForInputName(const ActionsDAGPtr & actions);
+    const ActionsDAG::Node * find(const String & name);
 
 private:
     ActionsDAGPtr actions;

src/Processors/QueryPlan/ExpressionStep.cpp

@@ -79,27 +79,15 @@ void ExpressionStep::updateOutputStream()
     if (!getDataStreamTraits().preserves_sorting)
         return;
 
-    FindOriginalNodeForOutputName original_node_finder(actions_dag);
+    FindAliasForInputName alias_finder(actions_dag);
     const auto & input_sort_description = getInputStreams().front().sort_description;
     for (size_t i = 0, s = input_sort_description.size(); i < s; ++i)
     {
-        const auto & desc = input_sort_description[i];
         String alias;
-        const auto & origin_column = desc.column_name;
-        for (const auto & column : output_stream->header)
-        {
-            const auto * original_node = original_node_finder.find(column.name);
-            if (original_node && original_node->result_name == origin_column)
-            {
-                alias = column.name;
-                break;
-            }
-        }
-
-        if (alias.empty())
-            return;
-
-        output_stream->sort_description[i].column_name = alias;
+        const auto & original_column = input_sort_description[i].column_name;
+        const auto * alias_node = alias_finder.find(original_column);
+        if (alias_node)
+            output_stream->sort_description[i].column_name = alias_node->result_name;
     }
 }
 

src/Processors/QueryPlan/FilterStep.cpp

@@ -109,27 +109,15 @@ void FilterStep::updateOutputStream()
     if (!getDataStreamTraits().preserves_sorting)
         return;
 
-    FindOriginalNodeForOutputName original_node_finder(actions_dag);
+    FindAliasForInputName alias_finder(actions_dag);
     const auto & input_sort_description = getInputStreams().front().sort_description;
     for (size_t i = 0, s = input_sort_description.size(); i < s; ++i)
     {
-        const auto & desc = input_sort_description[i];
         String alias;
-        const auto & origin_column = desc.column_name;
-        for (const auto & column : output_stream->header)
-        {
-            const auto * original_node = original_node_finder.find(column.name);
-            if (original_node && original_node->result_name == origin_column)
-            {
-                alias = column.name;
-                break;
-            }
-        }
-
-        if (alias.empty())
-            return;
-
-        output_stream->sort_description[i].column_name = alias;
+        const auto & original_column = input_sort_description[i].column_name;
+        const auto * alias_node = alias_finder.find(original_column);
+        if (alias_node)
+            output_stream->sort_description[i].column_name = alias_node->result_name;
     }
 }