Merge 6ea7994 into b92e8dc

kinash-varvara · web-flow · commit 1d89486acb2e · 2024-02-18T11:16:02.000Z
diff --git a/library/cpp/json/json_writer.h b/library/cpp/json/json_writer.h
@@ -29,6 +29,12 @@ namespace NJson {
             return *this;
         }
 
+        inline TJsonWriterConfig& SetNanAsString(bool v) noexcept {
+            WriteNanAsString = v;
+
+            return *this;
+        }
+
         inline TJsonWriterConfig& SetFormatOutput(bool v) noexcept {
             FormatOutput = v;
 
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.SetNanAsString(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());

Original file line number	Diff line number	Diff line change
`@@ -323,10 +323,13 @@ TUnboxedValue TryParseJsonDom(const TStringBuf json, const IValueBuilder* valueB`
`323`	`323`	`}`
`324`	`324`	`}`
`325`	`325`
`326`		`-TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity, bool encodeUtf8) {`
	`326`	`+TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity, bool encodeUtf8, bool writeNanAsString) {`
`327`	`327`	`TStringStream output;`
`328`	`328`	`TJsonWriterConfig config;`
	`329`	`+`
`329`	`330`	`config.SetFormatOutput(false);`
	`331`	`+ config.SetNanAsString(writeNanAsString);`
	`332`	`+`
`330`	`333`	`config.FloatToStringMode = EFloatToStringMode::PREC_AUTO;`
`331`	`334`	`TJsonWriter writer(&output, config);`
`332`	`335`	`if (skipMapEntity)`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,6 @@ bool IsValidJson(const TStringBuf json);`
`9`	`9`
`10`	`10`	`NUdf::TUnboxedValue TryParseJsonDom(const TStringBuf json, const NUdf::IValueBuilder* valueBuilder, bool decodeUtf8 = false);`
`11`	`11`
`12`		`-TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity = false, bool encodeUtf8 = false);`
	`12`	`+TString SerializeJsonDom(const NUdf::TUnboxedValuePod dom, bool skipMapEntity = false, bool encodeUtf8 = false, bool writeNanAsString = false);`
`13`	`13`
`14`	`14`	`}`
Original file line number	Diff line number	Diff line change
`@@ -114,6 +114,11 @@`
`114`	`114`	`"uri": "file://test.test_JsonWithUtf8_/results.txt"`
`115`	`115`	`}`
`116`	`116`	`],`
	`117`	`+ "test.test[JsonWithNanAsString]": [`
	`118`	`+ {`
	`119`	`+ "uri": "file://test.test_JsonWithNanAsString_/results.txt"`
	`120`	`+ }`
	`121`	`+ ],`
`117`	`122`	`"test.test[Lists]": [`
`118`	`123`	`{`
`119`	`124`	`"uri": "file://test.test_Lists_/results.txt"`