Fix nullif implementation (#3392)

ydb-platform · Apr 3, 2024 · 34ddb63 · 34ddb63
1 parent d42a420
commit 34ddb63
Show file tree

Hide file tree

Showing 10 changed files with 91 additions and 20 deletions.
diff --git a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -2625,6 +2625,28 @@ TExprNode::TPtr ExpandPgArrayOp(const TExprNode::TPtr& input, TExprContext& ctx)
         .Build();
 }
 
+TExprNode::TPtr ExpandPgNullIf(const TExprNode::TPtr& input, TExprContext& ctx) {
+    auto pred = ctx.Builder(input->Pos())
+        .Callable("Coalesce")
+            .Callable(0, "FromPg")
+                .Callable( 0, "PgOp")
+                    .Atom(0, "=")
+                    .Add(1, input->Child(0))
+                    .Add(2, input->Child(1))
+                .Seal()
+            .Seal()
+            .Callable( 1, "Bool")
+                .Atom(0, "false")
+            .Seal()
+        .Seal().Build();
+    return ctx.Builder(input->Pos())
+        .Callable("If")
+            .Add(0, pred)
+            .Callable(1, "Null").Seal()
+            .Add(2, input->Child(0))
+        .Seal().Build();
+}
+
 template <bool Flat, bool List>
 TExprNode::TPtr ExpandContainerIf(const TExprNode::TPtr& input, TExprContext& ctx) {
     YQL_CLOG(DEBUG, CorePeepHole) << "Expand " << input->Content();
@@ -7868,6 +7890,7 @@ struct TPeepHoleRules {
         {"Contains", &RewriteSearchByKeyForTypesMismatch<true>},
         {"ListHas", &ExpandListHas},
         {"PgAnyResolvedOp", &ExpandPgArrayOp},
+        {"PgNullIf", &ExpandPgNullIf},
         {"PgAllResolvedOp", &ExpandPgArrayOp},
         {"Map", &CleckClosureOnUpperLambdaOverList},
         {"OrderedMap", &CleckClosureOnUpperLambdaOverList},

diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -12283,6 +12283,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["FromPg"] = &FromPgWrapper;
         Functions["ToPg"] = &ToPgWrapper;
         Functions["PgClone"] = &PgCloneWrapper;
+        Functions["PgNullIf"] = &PgNullIfWrapper;
         ExtFunctions["PgAgg"] = &PgAggWrapper;
         ExtFunctions["PgAggWindowCall"] = &PgAggWrapper;
         ExtFunctions["PgCall"] = &PgCallWrapper;

diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -917,6 +917,55 @@ IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode:
     return IGraphTransformer::TStatus::Ok;
 }
 
+IGraphTransformer::TStatus PgNullIfWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+    Y_UNUSED(output);
+    if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    TVector<ui32> types(2);
+    bool replaced = false;
+
+    for (ui32 i = 0; i < 2; ++i) {
+        auto* item = input->Child(i);
+        auto type = item->GetTypeAnn();
+        ui32 argType;
+        bool convertToPg;
+        const auto pos = item->Pos();
+        if (!ExtractPgType(type, argType, convertToPg, pos, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (convertToPg) {
+            replaced = true;
+            input->ChildRef(i) = ctx.Expr.NewCallable(input->Child(2)->Pos(), "ToPg", { input->ChildPtr(i) });
+        }
+
+        types[i] = argType;
+    }
+
+    if (replaced) {
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    const NPg::TTypeDesc* commonType;
+    if (const auto issue = NPg::LookupCommonType(types,
+        [&input, &ctx](size_t i) {
+            return ctx.Expr.GetPosition(input->Child(i)->Pos());
+        }, commonType))
+    {
+        ctx.Expr.AddError(*issue);
+        return IGraphTransformer::TStatus::Error;
+    }
+    if (IsCastRequired(commonType->TypeId, types[0])) {
+        input->ChildRef(0) = WrapWithPgCast(std::move(input->ChildRef(0)), commonType->TypeId, ctx.Expr);
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    input->SetTypeAnn(ctx.Expr.MakeType<TPgExprType>(commonType->TypeId));
+    return IGraphTransformer::TStatus::Ok;
+}
+
 IGraphTransformer::TStatus PgQualifiedStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
     Y_UNUSED(output);
     if (!EnsureArgsCount(*input, 1, ctx.Expr)) {

diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.h b/ydb/library/yql/core/type_ann/type_ann_pg.h
@@ -28,6 +28,7 @@ IGraphTransformer::TStatus PgCloneWrapper(const TExprNode::TPtr& input, TExprNod
 IGraphTransformer::TStatus PgOpWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgArrayOpWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
+IGraphTransformer::TStatus PgNullIfWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgQualifiedStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgColumnRefWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -4395,10 +4395,7 @@ class TConverter : public IPGParseEvents {
         if (!lhs || !rhs) {
             return nullptr;
         }
-        auto pred = L(A("Coalesce"),
-                L(A("FromPg"), L(A("PgOp"), QA("="), lhs, rhs)),
-                L(A("Bool"), QA("false")));
-        return L(A("If"), pred, lhs, L(A("Null")));
+        return L(A("PgNullIf"), lhs, rhs);
     }
 
     TAstNode* ParseAExprIn(const A_Expr* value, const TExprSettings& settings) {

diff --git a/ydb/library/yql/tests/sql/dq_file/part14/canondata/result.json b/ydb/library/yql/tests/sql/dq_file/part14/canondata/result.json
@@ -2180,9 +2180,9 @@
     ],
     "test.test[pg-nullif-default.txt-Debug]": [
         {
-            "checksum": "a19676d0aca7eeb4d766d30dbaa027a9",
-            "size": 615,
-            "uri": "https://{canondata_backend}/1937027/9074da5ec3159ab717d6f0fee0639313448b4579/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql_patched"
+            "checksum": "15878810614467bf104f998fba1655bd",
+            "size": 737,
+            "uri": "https://{canondata_backend}/1775319/aeec36fe6eb29fa67191adf64023416d42730ad7/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql_patched"
         }
     ],
     "test.test[pg-nullif-default.txt-Plan]": [

diff --git a/ydb/library/yql/tests/sql/hybrid_file/part6/canondata/result.json b/ydb/library/yql/tests/sql/hybrid_file/part6/canondata/result.json
@@ -1709,9 +1709,9 @@
     ],
     "test.test[pg-nullif-default.txt-Debug]": [
         {
-            "checksum": "9144cd982be865aa47ca681f7608114e",
-            "size": 614,
-            "uri": "https://{canondata_backend}/1937001/a770b7e950bbaeaf08ef4bbb336b7e3683a914ce/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql_patched"
+            "checksum": "4dcb2b695a7bd270235ed841f2987d61",
+            "size": 736,
+            "uri": "https://{canondata_backend}/1814674/0638f81fbaf2b2ba2ce2e6de8b8a96988a67d749/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql_patched"
         }
     ],
     "test.test[pg-nullif-default.txt-Plan]": [

diff --git a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
@@ -11390,9 +11390,9 @@
     ],
     "test_sql2yql.test[pg-nullif]": [
         {
-            "checksum": "dae077581c71350f52e42fed360dc646",
-            "size": 725,
-            "uri": "https://{canondata_backend}/937458/4426696fc50768982e8cf8f63bc4bda5708da9f4/resource.tar.gz#test_sql2yql.test_pg-nullif_/sql.yql"
+            "checksum": "5ad161b9dc9cb7a9697cdbf6486d4967",
+            "size": 957,
+            "uri": "https://{canondata_backend}/1689644/9b14941c017976925995bc2ca0ac4cd046275eea/resource.tar.gz#test_sql2yql.test_pg-nullif_/sql.yql"
         }
     ],
     "test_sql2yql.test[pg-order_by_agg_extra_for_keys]": [

diff --git a/ydb/library/yql/tests/sql/suites/pg/nullif.sql b/ydb/library/yql/tests/sql/suites/pg/nullif.sql
@@ -1,2 +1,2 @@
 --!syntax_pg
-select nullif(1, 1) as x, nullif(2, 1) as y;
+select nullif(1, 1), nullif(1, 2.2), nullif(2, 1), nullif(1.2, '7'), nullif(1.2, '1.2');
diff --git a/ydb/library/yql/tests/sql/yt_native_file/part14/canondata/result.json b/ydb/library/yql/tests/sql/yt_native_file/part14/canondata/result.json
@@ -2120,9 +2120,9 @@
     ],
     "test.test[pg-nullif-default.txt-Debug]": [
         {
-            "checksum": "a772808658bb9749b940ccafd9cbd0c8",
-            "size": 554,
-            "uri": "https://{canondata_backend}/1599023/7a3b3fda17e3f5ea5a2102f6dc324ab02a3ba67c/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql"
+            "checksum": "e92ad4e8c9de59bdc528ffc9f7547611",
+            "size": 677,
+            "uri": "https://{canondata_backend}/1925842/1ad87feba078148d2bd94b42ca0bec2e4edbaee6/resource.tar.gz#test.test_pg-nullif-default.txt-Debug_/opt.yql"
         }
     ],
     "test.test[pg-nullif-default.txt-Plan]": [
@@ -2134,9 +2134,9 @@
     ],
     "test.test[pg-nullif-default.txt-Results]": [
         {
-            "checksum": "a04c967abb8d43ec8e21592c6f6a5817",
-            "size": 938,
-            "uri": "https://{canondata_backend}/1599023/7a3b3fda17e3f5ea5a2102f6dc324ab02a3ba67c/resource.tar.gz#test.test_pg-nullif-default.txt-Results_/results.txt"
+            "checksum": "d765d0c205914bea1a9f3c1f546d7368",
+            "size": 1829,
+            "uri": "https://{canondata_backend}/1925842/1ad87feba078148d2bd94b42ca0bec2e4edbaee6/resource.tar.gz#test.test_pg-nullif-default.txt-Results_/results.txt"
         }
     ],
     "test.test[pg-range_function_multi-default.txt-Debug]": [