Fix typing in PgAgg with EmitAggApply enabled

Fixed typing in PgAgg with EmitAggApply

3 files changed, 39 insertions, 47 deletions

diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
index 46395ff771e..17e464d591e 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -630,28 +630,44 @@ IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode:
         return IGraphTransformer::TStatus::Repeat;
     }
 
+    const NPg::TAggregateDesc* aggDescPtr;
     try {
-        const auto& aggDesc = NPg::LookupAggregation(TString(name), argTypes);
-        if (aggDesc.Kind != NPg::EAggKind::Normal) {
+        aggDescPtr = &NPg::LookupAggregation(TString(name), argTypes);
+        if (aggDescPtr->Kind != NPg::EAggKind::Normal) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
                 "Only normal aggregation supported"));
             return IGraphTransformer::TStatus::Error;
         }
-
-        ui32 resultType;
-        if (!aggDesc.FinalFuncId) {
-            resultType = aggDesc.TransTypeId;
-        } else {
-            resultType = NPg::LookupProc(aggDesc.FinalFuncId).ResultType;
-        }
-
-        auto result = ctx.Expr.MakeType<TPgExprType>(resultType);
-        input->SetTypeAnn(result);
-        return IGraphTransformer::TStatus::Ok;
     } catch (const yexception& e) {
         ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), e.what()));
         return IGraphTransformer::TStatus::Error;
     }
+    const NPg::TAggregateDesc& aggDesc = *aggDescPtr;
+
+    ui32 argIdx = overWindow ? 3 : 2;
+    for (ui32 i = 0; i < argTypes.size(); ++i, ++argIdx) {
+        if (IsCastRequired(argTypes[i], aggDesc.ArgTypes[i])) {
+            auto& argNode = input->ChildRef(argIdx);
+            argNode = WrapWithPgCast(std::move(argNode), aggDesc.ArgTypes[i], ctx);
+            needRetype = true;
+        }
+    }
+
+    if (needRetype) {
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    ui32 resultType;
+    if (!aggDesc.FinalFuncId) {
+        resultType = aggDesc.TransTypeId;
+    } else {
+        resultType = NPg::LookupProc(aggDesc.FinalFuncId).ResultType;
+    }
+
+    auto result = ctx.Expr.MakeType<TPgExprType>(resultType);
+    input->SetTypeAnn(result);
+
+    return IGraphTransformer::TStatus::Ok;
 }
 
 IGraphTransformer::TStatus PgQualifiedStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
@@ -1265,12 +1281,6 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
     }
 
     const NPg::TAggregateDesc& aggDesc = *aggDescPtr;
-    for (size_t i = 0; i < argTypes.size(); ++i) {
-        if (IsCastRequired(argTypes[i], aggDesc.ArgTypes[i])) {
-            auto& argNode = lambda->ChildRef(i+1);
-            argNode = WrapWithPgCast(std::move(argNode), aggDesc.ArgTypes[i], ctx);
-        }
-    }
     output = ExpandPgAggregationTraits(input->Pos(), aggDesc, onWindow, lambda, argTypes, itemType, ctx.Expr);
     return IGraphTransformer::TStatus::Repeat;
 }
diff --git a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
index 8b56a86c3ff..6141ef68c0e 100644
--- a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
+++ b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
@@ -10527,6 +10527,13 @@
             "uri": "https://storage.yandex-team.ru/get-devtools/1817427/5ebfc414781230f90b896d4ed143ea00cf4c7aaa/resource.tar.gz#test_sql2yql.test_pg-aggregate_distinct_/sql.yql"
         }
     ],
+    "test_sql2yql.test[pg-aggregate_emit_agg_apply]": [
+        {
+            "checksum": "6fe3e9102614378326fdef1ebdc3d9f7",
+            "size": 724,
+            "uri": "https://storage.yandex-team.ru/get-devtools/1599023/83bd1e70db2f3ff932bdff8520b19de955cc6ddf/resource.tar.gz#test_sql2yql.test_pg-aggregate_emit_agg_apply_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[pg-aggregate_factory]": [
         {
             "checksum": "9807c37c4603181699c3903e43504ae1",
@@ -15742,20 +15749,6 @@
             "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql2yql.test_view-file_eval_/sql.yql"
         }
     ],
-    "test_sql2yql.test[view-file_inner]": [
-        {
-            "checksum": "d423dd9f7dc961eef7b7d4c424058618",
-            "size": 1405,
-            "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql2yql.test_view-file_inner_/sql.yql"
-        }
-    ],
-    "test_sql2yql.test[view-file_outer]": [
-        {
-            "checksum": "34859678471993ed589f818c72aeacb1",
-            "size": 1405,
-            "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql2yql.test_view-file_outer_/sql.yql"
-        }
-    ],
     "test_sql2yql.test[view-init_view_after_eval]": [
         {
             "checksum": "db1af9f900a72e1f750beed31c5e898b",
@@ -30337,20 +30330,6 @@
             "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql_format.test_view-file_eval_/formatted.sql"
         }
     ],
-    "test_sql_format.test[view-file_inner]": [
-        {
-            "checksum": "8f012ee9373d675c75b1649b96e6ead0",
-            "size": 116,
-            "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql_format.test_view-file_inner_/formatted.sql"
-        }
-    ],
-    "test_sql_format.test[view-file_outer]": [
-        {
-            "checksum": "9bf795fe9976b5eb6aa0018f00f9dbf9",
-            "size": 116,
-            "uri": "https://storage.yandex-team.ru/get-devtools/1880306/50efd53ad53129cbd4f94db66527b230b24320d3/resource.tar.gz#test_sql_format.test_view-file_outer_/formatted.sql"
-        }
-    ],
     "test_sql_format.test[view-init_view_after_eval]": [
         {
             "checksum": "2b80cc480f3257c380fa17a433bf7d55",
diff --git a/ydb/library/yql/tests/sql/suites/pg/aggregate_emit_agg_apply.sql b/ydb/library/yql/tests/sql/suites/pg/aggregate_emit_agg_apply.sql
new file mode 100644
index 00000000000..f3f663340c5
--- /dev/null
+++ b/ydb/library/yql/tests/sql/suites/pg/aggregate_emit_agg_apply.sql
@@ -0,0 +1,3 @@
+--!syntax_pg
+set EmitAggApply=true;
+select string_agg(cast(x as text),',') from generate_series(1,2) as x