Map some YQL operators to PgOp

3 files changed, 54 insertions, 0 deletions

diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 81d50565c6..89388ddbc1 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -2256,6 +2256,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus AddWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("+", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         const bool checked = input->Content().StartsWith("Checked");
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
@@ -2344,6 +2348,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus SubWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("-", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         const bool checked = input->Content().StartsWith("Checked");
 
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -2431,6 +2439,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus MulWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("*", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         const bool checked = input->Content().StartsWith("Checked");
 
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -2505,6 +2517,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus DivWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("/", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         const bool checked = input->Content().StartsWith("Checked");
 
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -2574,6 +2590,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus ModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("%", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+        
         const bool checked = input->Content().StartsWith("Checked");
 
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -2784,6 +2804,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus PlusMinusWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp(input->Content() == "Plus" ? "+" : "-", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2962,6 +2986,10 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus ConcatWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (auto status = TryConvertToPgOp("||", input, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index 1c02aee73a..d916abd98d 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -5394,6 +5394,31 @@ bool HasContextFuncs(const TExprNode& input) {
     return needCtx;
 }
 
+IGraphTransformer::TStatus TryConvertToPgOp(TStringBuf op, const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
+    if (!EnsureMinArgsCount(*input, 1, ctx)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    bool hasPg = false;
+    for (const auto& child : input->Children()) {
+        if (!EnsureComputable(*child, ctx)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        hasPg = hasPg || child->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Pg;
+    }
+
+    if (!hasPg) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    TExprNode::TListType args;
+    args.push_back(ctx.NewAtom(input->Pos(), op));
+    args.insert(args.end(), input->Children().begin(), input->Children().end());
+    output = ctx.NewCallable(input->Pos(), "PgOp", std::move(args));
+    return IGraphTransformer::TStatus::Repeat;
+}
+
 bool EnsureBlockOrScalarType(const TExprNode& node, TExprContext& ctx) {
     if (HasError(node.GetTypeAnn(), ctx) || !node.GetTypeAnn()) {
         YQL_ENSURE(node.Type() == TExprNode::Lambda);
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index b3ce8d9666..05d4d4152b 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -302,6 +302,7 @@ std::optional<ui32> GetFieldPosition(const TStructExprType& structType, const TS
 
 bool ExtractPgType(const TTypeAnnotationNode* type, ui32& pgType, bool& convertToPg, TPositionHandle pos, TExprContext& ctx);
 bool HasContextFuncs(const TExprNode& input);
+IGraphTransformer::TStatus TryConvertToPgOp(TStringBuf op, const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx);
 
 bool EnsureBlockOrScalarType(const TExprNode& node, TExprContext& ctx);
 bool EnsureBlockOrScalarType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);