Merge pull request #12606 from vitlibar/improve-json-internal-interface

Improve JSON internal interface.

src/Functions/DummyJSONParser.h

@@ -17,6 +17,8 @@ struct DummyJSONParser
     class Array;
     class Object;
 
+    /// References an element in a JSON document, representing a JSON null, boolean, string, number,
+    /// array or object.
     class Element
     {
     public:
@@ -39,6 +41,7 @@ struct DummyJSONParser
         Object getObject() const;
     };
 
+    /// References an array in a JSON document.
     class Array
     {
     public:
@@ -46,10 +49,10 @@ struct DummyJSONParser
         {
         public:
             Element operator*() const { return {}; }
-            Iterator & operator ++() { return *this; }
+            Iterator & operator++() { return *this; }
-            Iterator operator ++(int) { return *this; }
+            Iterator operator++(int) { return *this; }
-            friend bool operator ==(const Iterator &, const Iterator &) { return true; }
+            friend bool operator==(const Iterator &, const Iterator &) { return true; }
-            friend bool operator !=(const Iterator &, const Iterator &) { return false; }
+            friend bool operator!=(const Iterator &, const Iterator &) { return false; }
         };
 
         Iterator begin() const { return {}; }
@@ -58,29 +61,40 @@ struct DummyJSONParser
         Element operator[](size_t) const { return {}; }
     };
 
+    using KeyValuePair = std::pair<std::string_view, Element>;
+
+    /// References an object in a JSON document.
     class Object
     {
     public:
-        using KeyValuePair = std::pair<std::string_view, Element>;
-
         class Iterator
         {
         public:
-            KeyValuePair operator *() const { return {}; }
+            KeyValuePair operator*() const { return {}; }
-            Iterator & operator ++() { return *this; }
+            Iterator & operator++() { return *this; }
-            Iterator operator ++(int) { return *this; }
+            Iterator operator++(int) { return *this; }
-            friend bool operator ==(const Iterator &, const Iterator &) { return true; }
+            friend bool operator==(const Iterator &, const Iterator &) { return true; }
-            friend bool operator !=(const Iterator &, const Iterator &) { return false; }
+            friend bool operator!=(const Iterator &, const Iterator &) { return false; }
         };
 
         Iterator begin() const { return {}; }
         Iterator end() const { return {}; }
         size_t size() const { return 0; }
-        KeyValuePair operator[](size_t) const { return {}; }
         bool find(const std::string_view &, Element &) const { return false; }
+
+#if 0
+        /// Optional: Provides access to an object's element by index.
+        KeyValuePair operator[](size_t) const { return {}; }
+#endif
     };
 
+    /// Parses a JSON document, returns the reference to its root element if succeeded.
     bool parse(const std::string_view &, Element &) { throw Exception{"Functions JSON* are not supported", ErrorCodes::NOT_IMPLEMENTED}; }
+
+#if 0
+    /// Optional: Allocates memory to parse JSON documents faster.
+    void reserve(size_t max_size);
+#endif
 };
 
 }

src/Functions/FunctionsJSON.h

@@ -139,6 +139,12 @@ private:
     BOOST_TTI_HAS_MEMBER_FUNCTION(reserve)
     BOOST_TTI_HAS_MEMBER_FUNCTION(prepare)
 
+    template <class T, class = void>
+    struct has_index_operator : std::false_type {};
+
+    template <class T>
+    struct has_index_operator<T, std::void_t<decltype(std::declval<T>()[0])>> : std::true_type {};
+
     /// Represents a move of a JSON iterator described by a single argument passed to a JSON function.
     /// For example, the call JSONExtractInt('{"a": "hello", "b": [-100, 200.0, 300]}', 'b', 1)
     /// contains two moves: {MoveType::ConstKey, "b"} and {MoveType::ConstIndex, 1}.
@@ -217,38 +223,32 @@ private:
         {
             auto array = element.getArray();
             if (index >= 0)
-            {
                 --index;
-                if (static_cast<size_t>(index) >= array.size())
+            else
-                    return false;
+                index += array.size();
-                element = array[index];
+
-                out_key = {};
+            if (static_cast<size_t>(index) >= array.size())
-                return true;
-            }
-            index += array.size();
-            if (index < 0)
                 return false;
             element = array[index];
             out_key = {};
             return true;
         }
 
-        if (element.isObject())
+        if constexpr (has_index_operator<typename JSONParser::Object>::value)
         {
-            auto object = element.getObject();
+            if (element.isObject())
-            if (index >= 0)
             {
-                --index;
+                auto object = element.getObject();
+                if (index >= 0)
+                    --index;
+                else
+                    index += object.size();
+
                 if (static_cast<size_t>(index) >= object.size())
                     return false;
                 std::tie(out_key, element) = object[index];
                 return true;
             }
-            index += object.size();
-            if (index < 0)
-                return false;
-            std::tie(out_key, element) = object[index];
-            return true;
         }
 
         return {};

src/Functions/RapidJSONParser.h

@@ -20,6 +20,8 @@ struct RapidJSONParser
     class Array;
     class Object;
 
+    /// References an element in a JSON document, representing a JSON null, boolean, string, number,
+    /// array or object.
     class Element
     {
     public:
@@ -47,6 +49,7 @@ struct RapidJSONParser
         const rapidjson::Value * ptr = nullptr;
     };
 
+    /// References an array in a JSON document.
     class Array
     {
     public:
@@ -67,17 +70,18 @@ struct RapidJSONParser
         ALWAYS_INLINE Iterator begin() const { return ptr->Begin(); }
         ALWAYS_INLINE Iterator end() const { return ptr->End(); }
         ALWAYS_INLINE size_t size() const { return ptr->Size(); }
-        ALWAYS_INLINE Element operator[](size_t index) const { return *(ptr->Begin() + index); }
+        ALWAYS_INLINE Element operator[](size_t index) const { assert(index < size()); return *(ptr->Begin() + index); }
 
     private:
         const rapidjson::Value * ptr = nullptr;
     };
 
+    using KeyValuePair = std::pair<std::string_view, Element>;
+
+    /// References an object in a JSON document.
     class Object
     {
     public:
-        using KeyValuePair = std::pair<std::string_view, Element>;
-
         class Iterator
         {
         public:
@@ -96,14 +100,7 @@ struct RapidJSONParser
         ALWAYS_INLINE Iterator end() const { return ptr->MemberEnd(); }
         ALWAYS_INLINE size_t size() const { return ptr->MemberCount(); }
 
-        ALWAYS_INLINE KeyValuePair operator[](size_t index) const
+        bool find(const std::string_view & key, Element & result) const
-        {
-            auto it = ptr->MemberBegin() + index;
-            std::string_view key{it->name.GetString(), it->name.GetStringLength()};
-            return KeyValuePair{key, it->value};
-        }
-
-        ALWAYS_INLINE bool find(const std::string_view & key, Element & result) const
         {
             auto it = ptr->FindMember(rapidjson::StringRef(key.data(), key.length()));
             if (it == ptr->MemberEnd())
@@ -113,10 +110,20 @@ struct RapidJSONParser
             return true;
         }
 
+        /// Optional: Provides access to an object's element by index.
+        ALWAYS_INLINE KeyValuePair operator[](size_t index) const
+        {
+            assert (index < size());
+            auto it = ptr->MemberBegin() + index;
+            std::string_view key{it->name.GetString(), it->name.GetStringLength()};
+            return {key, it->value};
+        }
+
     private:
         const rapidjson::Value * ptr = nullptr;
     };
 
+    /// Parses a JSON document, returns the reference to its root element if succeeded.
     bool parse(const std::string_view & json, Element & result)
     {
         rapidjson::MemoryStream ms(json.data(), json.size());
@@ -128,6 +135,11 @@ struct RapidJSONParser
         return true;
     }
 
+#if 0
+    /// Optional: Allocates memory to parse JSON documents faster.
+    void reserve(size_t max_size);
+#endif
+
 private:
     rapidjson::Document document;
 };

src/Functions/SimdJSONParser.h

@@ -25,6 +25,8 @@ struct SimdJSONParser
     class Array;
     class Object;
 
+    /// References an element in a JSON document, representing a JSON null, boolean, string, number,
+    /// array or object.
     class Element
     {
     public:
@@ -52,6 +54,7 @@ struct SimdJSONParser
         simdjson::dom::element element;
     };
 
+    /// References an array in a JSON document.
     class Array
     {
     public:
@@ -59,11 +62,11 @@ struct SimdJSONParser
         {
         public:
             ALWAYS_INLINE Iterator(const simdjson::dom::array::iterator & it_) : it(it_) {}
-            ALWAYS_INLINE Element operator *() const { return *it; }
+            ALWAYS_INLINE Element operator*() const { return *it; }
-            ALWAYS_INLINE Iterator & operator ++() { ++it; return *this; }
+            ALWAYS_INLINE Iterator & operator++() { ++it; return *this; }
-            ALWAYS_INLINE Iterator operator ++(int) { auto res = *this; ++it; return res; }
+            ALWAYS_INLINE Iterator operator++(int) { auto res = *this; ++it; return res; }
-            ALWAYS_INLINE friend bool operator !=(const Iterator & left, const Iterator & right) { return left.it != right.it; }
+            ALWAYS_INLINE friend bool operator!=(const Iterator & left, const Iterator & right) { return left.it != right.it; }
-            ALWAYS_INLINE friend bool operator ==(const Iterator & left, const Iterator & right) { return !(left != right); }
+            ALWAYS_INLINE friend bool operator==(const Iterator & left, const Iterator & right) { return !(left != right); }
         private:
             simdjson::dom::array::iterator it;
         };
@@ -72,26 +75,27 @@ struct SimdJSONParser
         ALWAYS_INLINE Iterator begin() const { return array.begin(); }
         ALWAYS_INLINE Iterator end() const { return array.end(); }
         ALWAYS_INLINE size_t size() const { return array.size(); }
-        ALWAYS_INLINE Element operator[](size_t index) const { return array.at(index).first; }
+        ALWAYS_INLINE Element operator[](size_t index) const { assert(index < size()); return array.at(index).first; }
 
     private:
         simdjson::dom::array array;
     };
 
+    using KeyValuePair = std::pair<std::string_view, Element>;
+
+    /// References an object in a JSON document.
     class Object
     {
     public:
-        using KeyValuePair = std::pair<std::string_view, Element>;
-
         class Iterator
         {
         public:
             ALWAYS_INLINE Iterator(const simdjson::dom::object::iterator & it_) : it(it_) {}
-            ALWAYS_INLINE KeyValuePair operator *() const { const auto & res = *it; return {res.key, res.value}; }
+            ALWAYS_INLINE KeyValuePair operator*() const { const auto & res = *it; return {res.key, res.value}; }
-            ALWAYS_INLINE Iterator & operator ++() { ++it; return *this; }
+            ALWAYS_INLINE Iterator & operator++() { ++it; return *this; }
-            ALWAYS_INLINE Iterator operator ++(int) { auto res = *this; ++it; return res; }
+            ALWAYS_INLINE Iterator operator++(int) { auto res = *this; ++it; return res; }
-            ALWAYS_INLINE friend bool operator !=(const Iterator & left, const Iterator & right) { return left.it != right.it; }
+            ALWAYS_INLINE friend bool operator!=(const Iterator & left, const Iterator & right) { return left.it != right.it; }
-            ALWAYS_INLINE friend bool operator ==(const Iterator & left, const Iterator & right) { return !(left != right); }
+            ALWAYS_INLINE friend bool operator==(const Iterator & left, const Iterator & right) { return !(left != right); }
         private:
             simdjson::dom::object::iterator it;
         };
@@ -101,15 +105,7 @@ struct SimdJSONParser
         ALWAYS_INLINE Iterator end() const { return object.end(); }
         ALWAYS_INLINE size_t size() const { return object.size(); }
 
-        KeyValuePair operator [](size_t index) const
+        bool find(const std::string_view & key, Element & result) const
-        {
-            Iterator it = begin();
-            while (index--)
-                ++it;
-            return *it;
-        }
-
-        ALWAYS_INLINE bool find(const std::string_view & key, Element & result) const
         {
             auto x = object.at_key(key);
             if (x.error())
@@ -119,17 +115,22 @@ struct SimdJSONParser
             return true;
         }
 
+        /// Optional: Provides access to an object's element by index.
+        KeyValuePair operator[](size_t index) const
+        {
+            assert(index < size());
+            auto it = object.begin();
+            while (index--)
+                ++it;
+            const auto & res = *it;
+            return {res.key, res.value};
+        }
+
     private:
         simdjson::dom::object object;
     };
 
-    void reserve(size_t max_size)
+    /// Parses a JSON document, returns the reference to its root element if succeeded.
-    {
-        if (parser.allocate(max_size) != simdjson::error_code::SUCCESS)
-            throw Exception{"Couldn't allocate " + std::to_string(max_size) + " bytes when parsing JSON",
-                            ErrorCodes::CANNOT_ALLOCATE_MEMORY};
-    }
-
     bool parse(const std::string_view & json, Element & result)
     {
         auto document = parser.parse(json.data(), json.size());
@@ -140,6 +141,14 @@ struct SimdJSONParser
         return true;
     }
 
+    /// Optional: Allocates memory to parse JSON documents faster.
+    void reserve(size_t max_size)
+    {
+        if (parser.allocate(max_size) != simdjson::error_code::SUCCESS)
+            throw Exception{"Couldn't allocate " + std::to_string(max_size) + " bytes when parsing JSON",
+                            ErrorCodes::CANNOT_ALLOCATE_MEMORY};
+    }
+
 private:
     simdjson::dom::parser parser;
 };