[yt provider] Fix YtReduce type annotation with aux sort columns

YQL-19373
commit_hash:8097d935a27fdc5e71ccee6855681514dcef892b

5 files changed, 60 insertions, 6 deletions

diff --git a/yql/essentials/tests/sql/sql2yql/canondata/result.json b/yql/essentials/tests/sql/sql2yql/canondata/result.json
index 6780fe3a7f..314d1e633d 100644
--- a/yql/essentials/tests/sql/sql2yql/canondata/result.json
+++ b/yql/essentials/tests/sql/sql2yql/canondata/result.json
@@ -11395,6 +11395,13 @@
             "uri": "https://{canondata_backend}/1937001/40d997912e417875462933d31ee934f62af7dd8a/resource.tar.gz#test_sql2yql.test_optimizers-pushdown_nonsep_over_aggregate_/sql.yql"
         }
     ],
+    "test_sql2yql.test[optimizers-reduce_with_aux_sort_column]": [
+        {
+            "checksum": "5d7c8ddf7ad11a2c7354ab8424a8031a",
+            "size": 2765,
+            "uri": "https://{canondata_backend}/1781765/cb24b9901fae16f9664ce3775954713b77cca79e/resource.tar.gz#test_sql2yql.test_optimizers-reduce_with_aux_sort_column_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[optimizers-sort_by_nonstrict_const]": [
         {
             "checksum": "169c512d23e47b59ae22f678661c4ac2",
@@ -28124,6 +28131,11 @@
             "uri": "file://test_sql_format.test_optimizers-pushdown_nonsep_over_aggregate_/formatted.sql"
         }
     ],
+    "test_sql_format.test[optimizers-reduce_with_aux_sort_column]": [
+        {
+            "uri": "file://test_sql_format.test_optimizers-reduce_with_aux_sort_column_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[optimizers-sort_by_nonstrict_const]": [
         {
             "uri": "file://test_sql_format.test_optimizers-sort_by_nonstrict_const_/formatted.sql"
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_optimizers-reduce_with_aux_sort_column_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_optimizers-reduce_with_aux_sort_column_/formatted.sql
new file mode 100644
index 0000000000..8758142d5f
--- /dev/null
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_optimizers-reduce_with_aux_sort_column_/formatted.sql
@@ -0,0 +1,28 @@
+USE plato;
+
+PRAGMA yt.JoinMergeTablesLimit = '10';
+
+$i = (
+    SELECT
+        a.key AS key,
+        a.subkey AS subkey,
+        b.value AS value
+    FROM
+        Input AS a
+    INNER JOIN
+        Input AS b
+    USING (key, subkey)
+    ASSUME ORDER BY
+        key,
+        subkey
+);
+
+INSERT INTO Output
+SELECT
+    AsTuple(key, subkey) AS k,
+    value || 'a' AS v
+FROM
+    $i
+ASSUME ORDER BY
+    k
+;
diff --git a/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.cfg b/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.cfg
new file mode 100644
index 0000000000..66737248b8
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.cfg
@@ -0,0 +1,2 @@
+in Input sorted.txt
+out Output output.txt
diff --git a/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.sql b/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.sql
new file mode 100644
index 0000000000..8cb88e34ff
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/optimizers/reduce_with_aux_sort_column.sql
@@ -0,0 +1,12 @@
+use plato;
+pragma yt.JoinMergeTablesLimit="10";
+
+$i = select a.key as key, a.subkey as subkey, b.value as value
+    from Input as a
+    inner join Input as b
+    using(key, subkey)
+assume order by key,subkey;
+
+insert into Output
+select AsTuple(key, subkey) as k, value || "a" as v from $i
+assume order by k;
diff --git a/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp b/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
index 4f53f48424..b422af1bf2 100644
--- a/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
@@ -250,7 +250,7 @@ private:
         return TStatus::Ok;
     }
 
-    static bool ValidateOutputType(const TTypeAnnotationNode* itemType, TPositionHandle positionHandle, TYtOutSection outTables,
+    static bool ValidateOutputType(const TTypeAnnotationNode* itemType, TPositionHandle positionHandle, const TYtOutSection& outTables,
         size_t beginIdx, size_t endIdx, bool useExtendedType, TExprContext& ctx)
     {
         YQL_ENSURE(beginIdx <= endIdx);
@@ -352,13 +352,13 @@ private:
         return true;
     }
 
-    static bool ValidateOutputType(const TTypeAnnotationNode* itemType, TPositionHandle positionHandle, TYtOutSection outTables,
+    static bool ValidateOutputType(const TTypeAnnotationNode* itemType, TPositionHandle positionHandle, const TYtOutSection& outTables,
         TExprContext& ctx, bool useExtendedType = false)
     {
         return ValidateOutputType(itemType, positionHandle, outTables, 0, outTables.Ref().ChildrenSize(), useExtendedType, ctx);
     }
 
-    static bool ValidateOutputType(const TExprNode& list, TYtOutSection outTables, TExprContext& ctx, bool useExtendedType = false) {
+    static bool ValidateOutputType(const TExprNode& list, const TYtOutSection& outTables, TExprContext& ctx, bool useExtendedType = false) {
         const TTypeAnnotationNode* itemType = GetSequenceItemType(list.Pos(), list.GetTypeAnn(), true, ctx);
         if (nullptr == itemType) {
             return false;
@@ -1144,7 +1144,7 @@ private:
             return TStatus::Error;
         }
 
-        if (!ValidateOutputType(*lambda, map.Output(), ctx, true)) {
+        if (!ValidateOutputType(*lambda, map.Output(), ctx, /*useExtendedType*/true)) {
             lambda->SetTypeAnn(nullptr);
             lambda->SetState(TExprNode::EState::Initial);
             return TStatus::Error;
@@ -1260,7 +1260,7 @@ private:
             return TStatus::Error;
         }
 
-        if (!ValidateOutputType(*lambda, reduce.Output(), ctx)) {
+        if (!ValidateOutputType(*lambda, reduce.Output(), ctx, /*useExtendedType*/true)) {
             lambda->SetTypeAnn(nullptr);
             lambda->SetState(TExprNode::EState::Initial);
             return TStatus::Error;
@@ -1638,7 +1638,7 @@ private:
             return TStatus::Error;
         }
 
-        if (!ValidateOutputType(lambda.Ref(), fill.Output(), ctx, true)) {
+        if (!ValidateOutputType(lambda.Ref(), fill.Output(), ctx, /*useExtendedType*/true)) {
             return TStatus::Error;
         }