optionally serialize double::inf in binary_json

commit_hash:837e450f67f5712de5717aa940433b0b36745315

3 files changed, 54 insertions, 19 deletions

diff --git a/yql/essentials/types/binary_json/ut/entry_ut.cpp b/yql/essentials/types/binary_json/ut/entry_ut.cpp
index 0f099ed59d9..247fc5034cf 100644
--- a/yql/essentials/types/binary_json/ut/entry_ut.cpp
+++ b/yql/essentials/types/binary_json/ut/entry_ut.cpp
@@ -17,6 +17,7 @@ public:
         UNIT_TEST(TestGetContainer);
         UNIT_TEST(TestGetString);
         UNIT_TEST(TestGetNumber);
+        UNIT_TEST(TestOutOfBounds);
     UNIT_TEST_SUITE_END();
 
     void TestGetType() {
@@ -93,6 +94,28 @@ public:
             UNIT_ASSERT_VALUES_EQUAL(container.GetElement(0).GetNumber(), testCase.second);
         }
     }
+
+    void TestOutOfBounds() {
+        const TVector<std::pair<TString, double>> testCases = {
+            { "1e100000000", std::numeric_limits<double>::infinity() },
+            { "-1e100000000", -std::numeric_limits<double>::infinity() },
+            { "1.797693135e+308", std::numeric_limits<double>::infinity() },
+            { "-1.797693135e+308", -std::numeric_limits<double>::infinity() },
+        };
+
+        for (const auto& testCase : testCases) {
+            UNIT_ASSERT(std::holds_alternative<TString>(SerializeToBinaryJson(testCase.first)));
+        }
+
+        for (const auto& testCase : testCases) {
+            const auto binaryJson = std::get<TBinaryJson>(SerializeToBinaryJson(testCase.first, true));
+            const auto reader = TBinaryJsonReader::Make(binaryJson);
+            const auto container = reader->GetRootCursor();
+
+            UNIT_ASSERT_VALUES_EQUAL(container.GetElement(0).GetNumber(), testCase.second);
+        }
+    }
 };
 
 UNIT_TEST_SUITE_REGISTRATION(TBinaryJsonEntryTest);
+
diff --git a/yql/essentials/types/binary_json/write.cpp b/yql/essentials/types/binary_json/write.cpp
index 28dffea9db2..39cd05949db 100644
--- a/yql/essentials/types/binary_json/write.cpp
+++ b/yql/essentials/types/binary_json/write.cpp
@@ -13,6 +13,8 @@
 #include <util/generic/set.h>
 #include <util/generic/stack.h>
 #include <util/generic/vector.h>
+#include <yql/essentials/minikql/dom/node.h>
+#include <yql/essentials/utils/parse_double.h>
 
 #include <cmath>
 
@@ -415,9 +417,8 @@ private:
  */
 class TBinaryJsonCallbacks : public TJsonCallbacks {
 public:
-    TBinaryJsonCallbacks(bool throwException)
-        : TJsonCallbacks(/* throwException */ throwException)
-    {
+    TBinaryJsonCallbacks(bool throwException, bool allowInf)
+        : TJsonCallbacks(/* throwException */ throwException), AllowInf(allowInf) {
     }
 
     bool OnNull() override {
@@ -445,7 +446,7 @@ public:
     }
 
     bool OnDouble(double value) override {
-        if (Y_UNLIKELY(std::isinf(value))) {
+        if (Y_UNLIKELY(std::isinf(value) && !AllowInf)) {
             if (ThrowException) {
                 ythrow yexception() << "JSON number is infinite";
             } else {
@@ -492,6 +493,7 @@ public:
 
 private:
     TJsonIndex Json;
+    bool AllowInf;
 };
 
 void DomToJsonIndex(const NUdf::TUnboxedValue& value, TBinaryJsonCallbacks& callbacks) {
@@ -573,8 +575,14 @@ template <typename TOnDemandValue>
             switch (value.get_number_type()) {
                 case simdjson::builtin::number_type::floating_point_number: {
                     double v;
-                    RETURN_IF_NOT_SUCCESS(value.get(v));
-                    callbacks.OnDouble(v);
+                    if (const auto& error = value.get(v); Y_UNLIKELY(error != simdjson::SUCCESS)) {
+                        if (!NYql::TryDoubleFromString((std::string_view)value.raw_json_token(), v)) {
+                            return error;
+                        }
+                    };
+                    if (Y_UNLIKELY(!callbacks.OnDouble(v))) {
+                        return simdjson::error_code::NUMBER_ERROR;
+                    }
                     break;
                 }
                 case simdjson::builtin::number_type::signed_integer: {
@@ -592,7 +600,9 @@ template <typename TOnDemandValue>
                 case simdjson::builtin::number_type::big_integer:
                     double v;
                     RETURN_IF_NOT_SUCCESS(value.get(v));
-                    callbacks.OnDouble(v);
+                    if (Y_UNLIKELY(!callbacks.OnDouble(v))) {
+                        return simdjson::error_code::NUMBER_ERROR;
+                    }
                     break;
             }
             break;
@@ -600,7 +610,7 @@ template <typename TOnDemandValue>
         case simdjson::ondemand::json_type::null: {
             auto is_null = value.is_null();
             RETURN_IF_NOT_SUCCESS(is_null.error());
-	    if (Y_UNLIKELY(!is_null.value_unsafe())) {
+            if (Y_UNLIKELY(!is_null.value_unsafe())) {
                 return simdjson::error_code::N_ATOM_ERROR;
             }
             callbacks.OnNull();
@@ -644,7 +654,7 @@ template <typename TOnDemandValue>
 }
 
 // unused, left for performance comparison
-[[maybe_unused]] [[nodiscard]] simdjson::error_code SimdJsonToJsonIndexImpl(const simdjson::dom::element& value, TBinaryJsonCallbacks& callbacks) {
+[[nodiscard]] [[maybe_unused]] simdjson::error_code SimdJsonToJsonIndexImpl(const simdjson::dom::element& value, TBinaryJsonCallbacks& callbacks) {
 #define RETURN_IF_NOT_SUCCESS(status)              \
     if (Y_UNLIKELY(status != simdjson::SUCCESS)) { \
         return status;                             \
@@ -715,9 +725,9 @@ template <typename TOnDemandValue>
 }
 }
 
-std::variant<TBinaryJson, TString> SerializeToBinaryJsonImpl(const TStringBuf json) {
+std::variant<TBinaryJson, TString> SerializeToBinaryJsonImpl(const TStringBuf json, bool allowInf) {
     std::variant<TBinaryJson, TString> res;
-    TBinaryJsonCallbacks callbacks(/* throwException */ false);
+    TBinaryJsonCallbacks callbacks(/* throwException */ false, allowInf);
     const simdjson::padded_string paddedJson(json);
     simdjson::ondemand::parser parser;
     try {
@@ -739,15 +749,16 @@ std::variant<TBinaryJson, TString> SerializeToBinaryJsonImpl(const TStringBuf js
     return res;
 }
 
-std::variant<TBinaryJson, TString> SerializeToBinaryJson(const TStringBuf json) {
-    return SerializeToBinaryJsonImpl(json);
+std::variant<TBinaryJson, TString> SerializeToBinaryJson(const TStringBuf json, bool allowInf) {
+    return SerializeToBinaryJsonImpl(json, allowInf);
 }
 
 TBinaryJson SerializeToBinaryJson(const NUdf::TUnboxedValue& value) {
-    TBinaryJsonCallbacks callbacks(/* throwException */ false);
+    TBinaryJsonCallbacks callbacks(/* throwException */ false, /* allowInf */ false);
     DomToJsonIndex(value, callbacks);
     TBinaryJsonSerializer serializer(std::move(callbacks).GetResult());
     return std::move(serializer).Serialize();
 }
 
 }
+
diff --git a/yql/essentials/types/binary_json/write.h b/yql/essentials/types/binary_json/write.h
index c4f8f89d2d2..61dea9bfb40 100644
--- a/yql/essentials/types/binary_json/write.h
+++ b/yql/essentials/types/binary_json/write.h
@@ -2,23 +2,24 @@
 
 #include "format.h"
 
-#include <yql/essentials/minikql/dom/node.h>
-
 #include <util/generic/maybe.h>
 
 #include <variant>
 
+namespace NYql::NUdf {
+class TUnboxedValue;
+};
+
 namespace NKikimr::NBinaryJson {
 
 /**
  * @brief Translates textual JSON into BinaryJson
  */
-std::variant<TBinaryJson, TString> SerializeToBinaryJson(const TStringBuf json);
+std::variant<TBinaryJson, TString> SerializeToBinaryJson(const TStringBuf json, bool allowInf = false);
 
 /**
  * @brief Translates DOM layout from `yql/library/dom` library into BinaryJson
  */
-TBinaryJson SerializeToBinaryJson(const NUdf::TUnboxedValue& value);
-
+TBinaryJson SerializeToBinaryJson(const NYql::NUdf::TUnboxedValue& value);
 }