YQL-13710 support context in YT combiner w/o keys

ref:72d162f1eca7ba871a3b42bac74b7b8b4c2dc682

4 files changed, 25 insertions, 12 deletions

diff --git a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
index eb63fbc11e4..6c1b345e185 100644
--- a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
+++ b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
@@ -2103,6 +2103,15 @@
             "Name": "TCoCurrentTzTimestamp",
             "Base": "TCoTzTimeBase",
             "Match": {"Type": "Callable", "Name": "CurrentTzTimestamp"}
+        },
+        {
+            "Name": "TCoWithContext",
+            "Base": "TCallable",
+            "Match": {"Type": "Callable", "Name": "WithContext"},
+            "Children": [
+                {"Index": 0, "Name": "Name", "Type": "TCoAtom"},
+                {"Index": 1, "Name": "Input", "Type": "TExprBase"}
+            ]
         }
     ]
 }
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index 223cd1e847b..420cd9554a4 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -2,6 +2,7 @@
 #include "yql_opt_proposed_by_data.h"
 #include "yql_opt_rewrite_io.h"
 #include "yql_opt_utils.h"
+#include "yql_expr_optimize.h"
 #include "yql_pg_utils.h"
 
 #include <ydb/library/yql/public/udf/udf_data_type.h>
@@ -4996,5 +4997,18 @@ bool ExtractPgType(const TTypeAnnotationNode* type, ui32& pgType, TPositionHandl
     }
 }
 
+bool HasContextFuncs(const TExprNode& input) {
+    bool needCtx = false;
+    VisitExpr(input, [&needCtx](const TExprNode& node) {
+        if (node.IsCallable("PgResolvedCallCtx")) {
+            needCtx = true;
+            return false;
+        }
+
+        return true;
+    });
+
+    return needCtx;
+}
 
 } // NYql
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index ad217ab5951..0561139ad26 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -281,5 +281,6 @@ std::optional<ui32> GetFieldPosition(const TStructExprType& structType, const TS
 
 bool IsCallableTypeHasStreams(const TCallableExprType* callableType);
 bool ExtractPgType(const TTypeAnnotationNode* type, ui32& pgType, TPositionHandle pos, TExprContext& ctx);
+bool HasContextFuncs(const TExprNode& input);
 
 }
diff --git a/ydb/library/yql/core/yql_opt_aggregate.cpp b/ydb/library/yql/core/yql_opt_aggregate.cpp
index fa2898c74cf..6b895bf7faa 100644
--- a/ydb/library/yql/core/yql_opt_aggregate.cpp
+++ b/ydb/library/yql/core/yql_opt_aggregate.cpp
@@ -2,7 +2,6 @@
 #include "yql_opt_utils.h"
 #include "yql_opt_window.h"
 #include "yql_expr_type_annotation.h"
-#include "yql_expr_optimize.h"
 
 namespace NYql {
 
@@ -36,17 +35,7 @@ TExprNode::TPtr ExpandAggregate(const TExprNode::TPtr& node, TExprContext& ctx,
         }
     }
 
-    bool needCtx = false;
-    VisitExpr(*aggregatedColumns, [&needCtx](const TExprNode& node) {
-        if (node.IsCallable("PgResolvedCallCtx")) {
-            needCtx = true;
-            return false;
-        }
-
-        return true;
-    });
-
-    auto contextLambda = needCtx ?
+    auto contextLambda = HasContextFuncs(*aggregatedColumns) ?
         ctx.Builder(node->Pos())
             .Lambda()
                 .Param("stream")