add WriteNanAsString option in Yson::SerializeJson (ydb-platform#1871)

GrigoriyPA · Feb 18, 2024 · a04d89b · a04d89b
1 parent b92e8dc
commit a04d89b
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 4 deletions.
diff --git a/ydb/library/yql/minikql/dom/json.cpp b/ydb/library/yql/minikql/dom/json.cpp
@@ -323,10 +323,13 @@ TUnboxedValue TryParseJsonDom(const TStringBuf json, const IValueBuilder* valueB
     }
 }
 
-TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity, bool encodeUtf8) {
+TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity, bool encodeUtf8, bool writeNanAsString) {
     TStringStream output;
     TJsonWriterConfig config;
+
     config.SetFormatOutput(false);
+    config.WriteNanAsString = writeNanAsString;
+
     config.FloatToStringMode = EFloatToStringMode::PREC_AUTO;
     TJsonWriter writer(&output, config);
     if (skipMapEntity)

diff --git a/ydb/library/yql/minikql/dom/json.h b/ydb/library/yql/minikql/dom/json.h
@@ -9,6 +9,6 @@ bool IsValidJson(const TStringBuf json);
 
 NUdf::TUnboxedValue TryParseJsonDom(const TStringBuf json, const NUdf::IValueBuilder* valueBuilder, bool decodeUtf8 = false);
 
-TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity = false, bool encodeUtf8 = false);
+TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity = false, bool encodeUtf8 = false, bool writeNanAsString = false);
 
 }
diff --git a/ydb/library/yql/minikql/dom/ut/yson_ut.cpp b/ydb/library/yql/minikql/dom/ut/yson_ut.cpp
@@ -1,5 +1,7 @@
 #include <ydb/library/yql/minikql/dom/yson.h>
 
+#include <ydb/library/yql/minikql/dom/json.h>
+
 #include <library/cpp/testing/unittest/registar.h>
 #include <ydb/library/yql/minikql/mkql_alloc.h>
 #include <ydb/library/yql/minikql/computation/mkql_computation_node_holders.h>
@@ -2059,4 +2061,27 @@ Y_UNIT_TEST_SUITE(TYsonTests) {
         const auto time = TInstant::Now() - t;
         Cerr << "Time is " << time << Endl;
     }
+
+    Y_UNIT_TEST(TestSerializeJsonNanInf) {
+        NMiniKQL::TScopedAlloc alloc(__LOCATION__);
+        NMiniKQL::TMemoryUsageInfo memInfo("Memory");
+        NMiniKQL::THolderFactory holderFactory(alloc.Ref(), memInfo, nullptr);
+        NMiniKQL::TDefaultValueBuilder builder(holderFactory);
+
+        constexpr char yson[] =
+        R"(
+        {
+            "Nan" = %nan;
+            "Inf" = %inf;
+            "NegInf" = %-inf
+        }
+        )";
+
+        TString expected(R"({"Inf":"inf","Nan":"nan","NegInf":"-inf"})");
+
+        const auto dom =  TryParseYsonDom(yson, &builder);
+        TString res = SerializeJsonDom(dom, false, true, true);
+
+        UNIT_ASSERT_EQUAL(expected, res);
+    }
 }
diff --git a/ydb/library/yql/udfs/common/yson2/test/canondata/result.json b/ydb/library/yql/udfs/common/yson2/test/canondata/result.json
@@ -114,6 +114,11 @@
             "uri": "file://test.test_JsonWithUtf8_/results.txt"
         }
     ],
+    "test.test[JsonWithNanAsString]": [
+        {
+            "uri": "file://test.test_JsonWithNanAsString_/results.txt"
+        }
+    ],
     "test.test[Lists]": [
         {
             "uri": "file://test.test_Lists_/results.txt"

diff --git a/ydb/library/yql/udfs/common/yson2/test/canondata/test.test_JsonWithNanAsString_/results.txt b/ydb/library/yql/udfs/common/yson2/test/canondata/test.test_JsonWithNanAsString_/results.txt
@@ -0,0 +1,59 @@
+[
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "Json"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "Json"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column2";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "Json"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "\"nan\""
+                        ];
+                        [
+                            "\"inf\""
+                        ];
+                        [
+                            "\"-inf\""
+                        ]
+                    ]
+                ]
+            }
+        ]
+    }
+]
diff --git a/ydb/library/yql/udfs/common/yson2/test/cases/JsonWithNanAsString.sql b/ydb/library/yql/udfs/common/yson2/test/cases/JsonWithNanAsString.sql
@@ -0,0 +1,8 @@
+$src =  Yson::From(0./0.);  --  nan
+$src1 = Yson::From(1./0.);  --  inf
+$src2 = Yson::From(-1./0.); -- -inf
+
+SELECT 
+    Yson::SerializeJson($src,  true AS WriteNanAsString),
+    Yson::SerializeJson($src1, true AS WriteNanAsString),
+    Yson::SerializeJson($src2, true AS WriteNanAsString)
diff --git a/ydb/library/yql/udfs/common/yson2/yson2_udf.cpp b/ydb/library/yql/udfs/common/yson2/yson2_udf.cpp
@@ -670,9 +670,10 @@ SIMPLE_STRICT_UDF(TSerializePretty, TYson(TAutoMap<TNodeResource>)) {
 
 constexpr char SkipMapEntity[] = "SkipMapEntity";
 constexpr char EncodeUtf8[] = "EncodeUtf8";
+constexpr char WriteNanAsString[] = "WriteNanAsString";
 
-SIMPLE_UDF_WITH_OPTIONAL_ARGS(TSerializeJson, TOptional<TJson>(TAutoMap<TNodeResource>, TOptional<TOptionsResource>, TNamedArg<bool, SkipMapEntity>, TNamedArg<bool, EncodeUtf8>), 3) try {
-    return valueBuilder->NewString(SerializeJsonDom(args[0], args[2].GetOrDefault(false), args[3].GetOrDefault(false)));
+SIMPLE_UDF_WITH_OPTIONAL_ARGS(TSerializeJson, TOptional<TJson>(TAutoMap<TNodeResource>, TOptional<TOptionsResource>, TNamedArg<bool, SkipMapEntity>, TNamedArg<bool, EncodeUtf8>, TNamedArg<bool, WriteNanAsString>), 4) try {
+    return valueBuilder->NewString(SerializeJsonDom(args[0], args[2].GetOrDefault(false), args[3].GetOrDefault(false), args[4].GetOrDefault(false)));
 } catch (const std::exception& e) {
     if (ParseOptions(args[1]).Strict) {
         UdfTerminate((::TStringBuilder() << valueBuilder->WithCalleePosition(GetPos()) << " " << e.what()).data());