KIKIMR-14495: new opt rule (drop Unordered over DqCnUnionAll)

ref:90743d930ed7f44953d8903ec87bf571f84dd133

2 files changed, 25 insertions, 0 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
index ee2f1fe6e6..66bc0482c0 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
@@ -57,6 +57,7 @@ public:
         AddHandler(0, &TKqlUpsertRowsIndex::Match, HNDL(BuildUpsertIndexStages));
         AddHandler(0, &TKqlInsertRowsIndex::Match, HNDL(BuildInsertIndexStages));
         AddHandler(0, &TKqlDeleteRowsIndex::Match, HNDL(BuildDeleteIndexStages));
+        AddHandler(0, &TCoUnorderedBase::Match, HNDL(DropUnordered));
         AddHandler(0, &TDqStage::Match, HNDL(FloatUpStage));
         AddHandler(0, &TCoHasItems::Match, HNDL(BuildHasItems));
         AddHandler(0, &TCoToOptional::Match, HNDL(BuildScalarPrecompute));
@@ -293,6 +294,15 @@ protected:
         return output;
     }
 
+    TMaybeNode<TExprBase> DropUnordered(TExprBase node, TExprContext& ctx) {
+        TExprBase output = node;
+        if (node.Maybe<TCoUnorderedBase>().Input().Maybe<TDqCnUnionAll>()) {
+            output = node.Cast<TCoUnorderedBase>().Input();
+        }
+        DumpAppliedRule("DropUnordered", node.Ptr(), output.Ptr(), ctx);
+        return output;
+    }
+
     TMaybeNode<TExprBase> FloatUpStage(TExprBase node, TExprContext& ctx) {
         TExprBase output = KqpFloatUpStage(node, ctx);
         DumpAppliedRule("FloatUpStage", node.Ptr(), output.Ptr(), ctx);
diff --git a/ydb/core/kqp/ut/kqp_ne_ut.cpp b/ydb/core/kqp/ut/kqp_ne_ut.cpp
index 1043baef5e..1c1f31acaa 100644
--- a/ydb/core/kqp/ut/kqp_ne_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_ut.cpp
@@ -3002,6 +3002,21 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
             UNIT_ASSERT(read["columns"].GetArraySafe().size() <= 2);
         }
     }
+
+    Y_UNIT_TEST(OrderedScalarContext) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExecuteDataQuery(R"(
+            PRAGMA kikimr.UseNewEngine = 'true';
+
+            $max_key = (SELECT Key FROM `/Root/KeyValue` ORDER BY Key DESC LIMIT 1);
+            SELECT $max_key ?? 0
+        )", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        CompareYson(R"([[2u]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
 }
 
 } // namespace NKikimr::NKqp