diff --git a/src/Processors/QueryPlan/Optimizations/liftUpFunctions.cpp b/src/Processors/QueryPlan/Optimizations/liftUpFunctions.cpp
index 80b82d989dd..32918f3e5a2 100644
--- a/src/Processors/QueryPlan/Optimizations/liftUpFunctions.cpp
+++ b/src/Processors/QueryPlan/Optimizations/liftUpFunctions.cpp
@@ -15,28 +15,12 @@ namespace ErrorCodes
 namespace
 {
 
-void swapSortingAndUnneededCalculations(DB::QueryPlan::Node * parent_node, DB::ActionsDAGPtr && unneeded_for_sorting)
+const DB::DataStream & getChildOutputStream(DB::QueryPlan::Node & node)
 {
-    DB::QueryPlan::Node * child_node = parent_node->children.front();
-
-    auto & parent_step = parent_node->step;
-    auto & child_step = child_node->step;
-    auto * sorting_step = typeid_cast<DB::SortingStep *>(parent_step.get());
-
-    // Sorting -> Expression
-    std::swap(parent_step, child_step);
-    // Expression -> Sorting
-
-    if (child_node->children.size() != 1)
-        throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "SortingStep is expected to have only one input stream.");
-    sorting_step->updateInputStream(child_node->children.front()->step->getOutputStream());
-    auto input_header = sorting_step->getInputStreams().front().header;
-    sorting_step->updateOutputStream(std::move(input_header));
-
-    auto description = parent_node->step->getStepDescription();
-    parent_step = std::make_unique<DB::ExpressionStep>(child_step->getOutputStream(), std::move(unneeded_for_sorting));
-    parent_step->setStepDescription(description + " [lifted up part]");
-    // UnneededCalculations -> Sorting
+    if (node.children.size() != 1)
+        throw DB::Exception(
+            DB::ErrorCodes::LOGICAL_ERROR, "Node \"{}\" is expected to have only one child.", node.step->getStepDescription());
+    return node.children.front()->step->getOutputStream();
 }
 
 }
@@ -68,20 +52,26 @@ size_t tryExecuteFunctionsAfterSorting(QueryPlan::Node * parent_node, QueryPlan:
     if (unneeded_for_sorting->trivial())
         return 0;
 
-    if (child_node->children.size() != 1)
-        throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "ExpressionStep is expected to have only one input stream.");
-
     // Sorting (parent_node) -> Expression (child_node)
     auto & node_with_needed = nodes.emplace_back();
     std::swap(node_with_needed.children, child_node->children);
     child_node->children = {&node_with_needed};
-    node_with_needed.step
-        = std::make_unique<ExpressionStep>(node_with_needed.children.front()->step->getOutputStream(), std::move(needed_for_sorting));
-    node_with_needed.step->setStepDescription(child_step->getStepDescription());
 
+    node_with_needed.step = std::make_unique<ExpressionStep>(getChildOutputStream(node_with_needed), std::move(needed_for_sorting));
+    node_with_needed.step->setStepDescription(child_step->getStepDescription());
     // Sorting (parent_node) -> so far the origin Expression (child_node) -> NeededCalculations (node_with_needed)
-    swapSortingAndUnneededCalculations(parent_node, std::move(unneeded_for_sorting));
-    // UneededCalculations (child_node) -> Sorting (parent_node) -> NeededCalculations (node_with_needed)
+
+    std::swap(parent_step, child_step);
+    // so far the origin Expression (parent_node) -> Sorting (child_node) -> NeededCalculations (node_with_needed)
+
+    sorting_step->updateInputStream(getChildOutputStream(*child_node));
+    auto input_header = sorting_step->getInputStreams().at(0).header;
+    sorting_step->updateOutputStream(std::move(input_header));
+
+    auto description = parent_step->getStepDescription();
+    parent_step = std::make_unique<DB::ExpressionStep>(child_step->getOutputStream(), std::move(unneeded_for_sorting));
+    parent_step->setStepDescription(description + " [lifted up part]");
+    // UneededCalculations (parent_node) -> Sorting (child_node) -> NeededCalculations (node_with_needed)
 
     return 3;
 }
diff --git a/src/Processors/QueryPlan/SortingStep.cpp b/src/Processors/QueryPlan/SortingStep.cpp
index efefbad0ded..859c9fd9e19 100644
--- a/src/Processors/QueryPlan/SortingStep.cpp
+++ b/src/Processors/QueryPlan/SortingStep.cpp
@@ -97,7 +97,7 @@ void SortingStep::updateInputStream(DataStream input_stream)
 
 void SortingStep::updateOutputStream(Block result_header)
 {
-    output_stream = createOutputStream(input_streams.front(), std::move(result_header), getDataStreamTraits());
+    output_stream = createOutputStream(input_streams.at(0), std::move(result_header), getDataStreamTraits());
     updateDistinctColumns(output_stream->header, output_stream->distinct_columns);
 }