diff --git a/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp b/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp index 10b8abc11886..1b1f4733a2b9 100644 --- a/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp +++ b/ydb/core/kqp/opt/rbo/kqp_rewrite_select.cpp @@ -343,6 +343,24 @@ TExprNode::TPtr ReplacePgOps(TExprNode::TPtr input, TExprContext &ctx) { } } +TVector GetSortDependencies(TExprNode::TPtr sort) { + TVector result; + + for (auto sortItem : sort->Child(1)->Children()) { + auto sortLambda = sortItem->Child(1); + TVector 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 sortElements; @@ -812,6 +830,7 @@ TExprNode::TPtr RewriteSelect(const TExprNode::TPtr &node, TExprContext &ctx, co finalColumnOrder.clear(); THashMap aggProjectionMap; + TVector finalProjection; for (auto resultItem : result->Child(1)->Children()) { auto column = resultItem->Child(0); @@ -909,6 +928,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(ctx, node->Pos()) + .Input(resultExpr) + .Variable().Value(iu.GetFullName()).Build() + .From().Value(iu.GetFullName()).Build() + .Done().Ptr()); + // clang-format on + } + } } // clang-format off @@ -921,11 +959,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 projectElements; + + for (auto c : finalProjection) { + // clang-format off + projectElements.push_back(Build(ctx, node->Pos()) + .Input(setItemPtr) + .Variable().Value(c).Build() + .From().Value(c).Build() + .Done().Ptr()); + // clang-format on + } + + // clang-format off + setItemPtr = Build(ctx, node->Pos()) + .Input(setItemPtr) + .MapElements() + .Add(projectElements) + .Build() + .Project() + .Value("true") + .Build() + .Done().Ptr(); + // clang-format on } setItemsResults.push_back(setItemPtr); diff --git a/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp b/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp index 94b223c8b2e1..daea44626223 100644 --- a/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp +++ b/ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp @@ -820,6 +820,18 @@ Y_UNIT_TEST_SUITE(KqpRboPg) { } std::vector 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/"; @@ -843,6 +855,8 @@ Y_UNIT_TEST_SUITE(KqpRboPg) { }; std::vector 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"]])"