[NEW RBO] Fixed order by support (#29359)

Co-authored-by: Pavel Velikhov <[email protected]>

Author: pavelvelikhov

ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp

@@ -339,6 +339,24 @@ TExprNode::TPtr ReplacePgOps(TExprNode::TPtr input, TExprContext &ctx) {
     }
 }
 
+TVector<TInfoUnit> GetSortDependencies(TExprNode::TPtr sort) {
+    TVector<TInfoUnit> result;
+
+    for (auto sortItem : sort->Child(1)->Children()) {   
+        auto sortLambda = sortItem->Child(1);
+        TVector<TInfoUnit> lambdaMembers;
+        
+        GetAllMembers(sortLambda, lambdaMembers);
+        for (auto m: lambdaMembers) {
+            if (std::find(result.begin(), result.end(), m) == result.end()) {
+                result.push_back(m);
+            }
+        }
+    }
+
+    return result; 
+}
+
 TExprNode::TPtr BuildSort(TExprNode::TPtr input, TExprNode::TPtr sort, TExprContext &ctx) {
     TVector<TExprNode::TPtr> sortElements;
 
@@ -807,6 +825,7 @@ TExprNode::TPtr RewriteSelect(const TExprNode::TPtr &node, TExprContext &ctx, co
 
         finalColumnOrder.clear();
         THashMap<TString, TExprNode::TPtr> aggProjectionMap;
+        TVector<TString> finalProjection;
 
         for (auto resultItem : result->Child(1)->Children()) {
             auto column = resultItem->Child(0);
@@ -904,6 +923,25 @@ TExprNode::TPtr RewriteSelect(const TExprNode::TPtr &node, TExprContext &ctx, co
                 .Lambda(lambda)
             .Done().Ptr());
             // clang-format on
+            
+            finalProjection.push_back(columnName);
+        }
+
+        // Sort clause may contain extra columns that we need to keep in the projection in order for sort to work
+        auto sort = GetSetting(setItem->Tail(), "sort");
+        if (sort) {
+            auto sortDependencies = GetSortDependencies(sort);
+            for (auto iu : sortDependencies) {
+                if (std::find(finalProjection.begin(), finalProjection.end(), iu.GetFullName()) == finalProjection.end()) {
+                    // clang-format off
+                    resultElements.push_back(Build<TKqpOpMapElementRename>(ctx, node->Pos())
+                        .Input(resultExpr)
+                        .Variable().Value(iu.GetFullName()).Build()
+                        .From().Value(iu.GetFullName()).Build()
+                    .Done().Ptr());
+                // clang-format on
+                }
+            }
         }
 
         // clang-format off
@@ -916,11 +954,34 @@ TExprNode::TPtr RewriteSelect(const TExprNode::TPtr &node, TExprContext &ctx, co
                 .Value("true")
             .Build()
         .Done().Ptr();
-        // clang-format onto
+        // clang-format on
 
-        auto sort = GetSetting(setItem->Tail(), "sort");
         if (sort) {
             setItemPtr = BuildSort(setItemPtr, sort, ctx);
+
+            TVector<TExprNode::TPtr> projectElements;
+
+            for (auto c : finalProjection) {
+                // clang-format off
+                projectElements.push_back(Build<TKqpOpMapElementRename>(ctx, node->Pos())
+                    .Input(setItemPtr)
+                    .Variable().Value(c).Build()
+                    .From().Value(c).Build()
+                .Done().Ptr());
+                // clang-format on
+            }
+
+            // clang-format off
+            setItemPtr = Build<TKqpOpMap>(ctx, node->Pos())
+                .Input(setItemPtr)
+                .MapElements()
+                    .Add(projectElements)
+                .Build()
+                .Project()
+                    .Value("true")
+                .Build()
+            .Done().Ptr();
+            // clang-format on
         }
 
         setItemsResults.push_back(setItemPtr);

ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp

@@ -846,6 +846,18 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         }
 
         std::vector<std::string> queries = {
+            R"(
+                --!syntax_pg
+                SET TablePathPrefix = "/Root/";
+                SELECT a FROM t1
+                ORDER BY a+1 DESC;
+            )",
+            R"(
+                --!syntax_pg
+                SET TablePathPrefix = "/Root/";
+                SELECT a FROM t1
+                ORDER BY a+1 DESC, c ASC;
+            )",
             R"(
                 --!syntax_pg
                 SET TablePathPrefix = "/Root/";
@@ -869,6 +881,8 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         };
 
         std::vector<std::string> results = {
+            R"([["3"];["2"];["1"];["0"]])",
+            R"([["3"];["2"];["1"];["0"]])",
             R"([["3"];["2"];["1"];["0"]])",
             R"([["3";"4"];["2";"3"];["1";"2"];["0";"1"]])",
             R"([["3"];["2"];["2"];["1"];["1"];["0"];["0"]])"