be/src/util/json/simd_json_parser.h

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/Common/JSONParsers/SimdJSONParser.h
// and modified by Doris

#pragma once

#include <simdjson.h>

#include <cassert>
#include <string>
#include <string_view>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

#include "core/types.h"
#include "util/string_parser.hpp"

namespace doris {

/// This class can be used as an argument for the template class FunctionJSON.
/// It provides ability to parse JSONs using simdjson library.
class SimdJSONParser {
    struct Node {
        enum class Type {
            INT64,
            UINT64,
            INT128,
            BIG_INTEGER,
            DOUBLE,
            STRING,
            ARRAY,
            OBJECT,
            BOOL,
            NULL_VALUE,
        };

        Type type = Type::NULL_VALUE;
        Int64 int64_value = 0;
        UInt64 uint64_value = 0;
        Int128 int128_value = 0;
        double double_value = 0;
        bool bool_value = false;
        std::string string_value;
        std::string raw_number;
        std::vector<Node> array_values;
        std::vector<std::string> object_keys;
        std::vector<Node> object_values;
    };

public:
    class Array;
    class Object;
    /// References an element in a JSON document, representing a JSON null, boolean, string, number,
    /// array or object.
    class Element {
    public:
        ALWAYS_INLINE Element() {} /// NOLINT
        ALWAYS_INLINE explicit Element(const simdjson::dom::element& element_)
                : dom_element(element_) {}                                 /// NOLINT
        ALWAYS_INLINE explicit Element(const Node* node_) : node(node_) {} /// NOLINT
        ALWAYS_INLINE bool isInt64() const {
            return node ? node->type == Node::Type::INT64
                        : dom_element.type() == simdjson::dom::element_type::INT64;
        }
        ALWAYS_INLINE bool isUInt64() const {
            return node ? node->type == Node::Type::UINT64
                        : dom_element.type() == simdjson::dom::element_type::UINT64;
        }
        ALWAYS_INLINE bool isInt128() const { return node && node->type == Node::Type::INT128; }
        ALWAYS_INLINE bool isBigInteger() const {
            return node && node->type == Node::Type::BIG_INTEGER;
        }
        ALWAYS_INLINE bool isDouble() const {
            return node ? node->type == Node::Type::DOUBLE
                        : dom_element.type() == simdjson::dom::element_type::DOUBLE;
        }
        ALWAYS_INLINE bool isString() const {
            return node ? node->type == Node::Type::STRING
                        : dom_element.type() == simdjson::dom::element_type::STRING;
        }
        ALWAYS_INLINE bool isArray() const {
            return node ? node->type == Node::Type::ARRAY
                        : dom_element.type() == simdjson::dom::element_type::ARRAY;
        }
        ALWAYS_INLINE bool isObject() const {
            return node ? node->type == Node::Type::OBJECT
                        : dom_element.type() == simdjson::dom::element_type::OBJECT;
        }
        ALWAYS_INLINE bool isBool() const {
            return node ? node->type == Node::Type::BOOL
                        : dom_element.type() == simdjson::dom::element_type::BOOLEAN;
        }
        ALWAYS_INLINE bool isNull() const {
            return node ? node->type == Node::Type::NULL_VALUE
                        : dom_element.type() == simdjson::dom::element_type::NULL_VALUE;
        }
        ALWAYS_INLINE Int64 getInt64() const {
            return node ? node->int64_value : dom_element.get_int64().value_unsafe();
        }
        ALWAYS_INLINE double getDouble() const {
            return node ? node->double_value : dom_element.get_double().value_unsafe();
        }
        ALWAYS_INLINE bool getBool() const {
            return node ? node->bool_value : dom_element.get_bool().value_unsafe();
        }
        ALWAYS_INLINE std::string_view getString() const {
            return node ? std::string_view(node->string_value)
                        : dom_element.get_string().value_unsafe();
        }
        ALWAYS_INLINE UInt64 getUInt64() const {
            return node ? node->uint64_value : dom_element.get_uint64().value_unsafe();
        }
        ALWAYS_INLINE Int128 getInt128() const {
            assert(node != nullptr);
            return node->int128_value;
        }
        ALWAYS_INLINE std::string_view getRawNumber() const {
            assert(node != nullptr);
            return node->raw_number;
        }
        ALWAYS_INLINE Array getArray() const;
        ALWAYS_INLINE Object getObject() const;

    private:
        simdjson::dom::element dom_element;
        const Node* node = nullptr;
    };
    /// References an array in a JSON document.
    class Array {
    public:
        class Iterator {
        public:
            using DomIterator = simdjson::dom::array::iterator;
            using NodeIterator = std::vector<Node>::const_iterator;
            ALWAYS_INLINE explicit Iterator(const DomIterator& it_) : it(it_) {} /// NOLINT
            ALWAYS_INLINE explicit Iterator(NodeIterator it_) : it(it_) {}       /// NOLINT
            ALWAYS_INLINE Element operator*() const {
                if (const auto* node_it = std::get_if<NodeIterator>(&it)) {
                    return Element(&**node_it);
                }
                return Element(*std::get<DomIterator>(it));
            }
            ALWAYS_INLINE Iterator& operator++() {
                if (auto* node_it = std::get_if<NodeIterator>(&it)) {
                    ++(*node_it);
                } else {
                    ++std::get<DomIterator>(it);
                }
                return *this;
            }
            ALWAYS_INLINE friend bool operator!=(const Iterator& left, const Iterator& right) {
                if (const auto* left_node_it = std::get_if<NodeIterator>(&left.it)) {
                    return *left_node_it != std::get<NodeIterator>(right.it);
                }
                return std::get<DomIterator>(left.it) != std::get<DomIterator>(right.it);
            }

        private:
            std::variant<DomIterator, NodeIterator> it;
        };
        ALWAYS_INLINE explicit Array(const simdjson::dom::array& array_)
                : dom_array(array_) {}                                                   /// NOLINT
        ALWAYS_INLINE explicit Array(const std::vector<Node>* array_) : array(array_) {} /// NOLINT
        ALWAYS_INLINE Iterator begin() const {
            return array ? Iterator(array->begin()) : Iterator(dom_array.begin());
        }
        ALWAYS_INLINE Iterator end() const {
            return array ? Iterator(array->end()) : Iterator(dom_array.end());
        }
        ALWAYS_INLINE size_t size() const { return array ? array->size() : dom_array.size(); }
        ALWAYS_INLINE Element operator[](size_t index) const {
            assert(index < size());
            return array ? Element(&(*array)[index]) : Element(dom_array.at(index).value_unsafe());
        }

    private:
        simdjson::dom::array dom_array;
        const std::vector<Node>* array = nullptr;
    };
    using KeyValuePair = std::pair<std::string_view, Element>;
    /// References an object in a JSON document.
    class Object {
    public:
        class Iterator {
        public:
            using DomIterator = simdjson::dom::object::iterator;
            ALWAYS_INLINE explicit Iterator(const std::vector<std::string>* keys_,
                                            const std::vector<Node>* values_, size_t index_)
                    : it(index_), keys(keys_), values(values_) {}                /// NOLINT
            ALWAYS_INLINE explicit Iterator(const DomIterator& it_) : it(it_) {} /// NOLINT
            ALWAYS_INLINE KeyValuePair operator*() const {
                if (const auto* index = std::get_if<size_t>(&it)) {
                    return {(*keys)[*index], Element(&(*values)[*index])};
                }
                const auto& res = *std::get<DomIterator>(it);
                return {res.key, Element(res.value)};
            }
            ALWAYS_INLINE Iterator& operator++() {
                if (auto* index = std::get_if<size_t>(&it)) {
                    ++(*index);
                } else {
                    ++std::get<DomIterator>(it);
                }
                return *this;
            }
            ALWAYS_INLINE Iterator operator++(int) {
                auto res = *this;
                ++(*this);
                return res;
            } /// NOLINT
            ALWAYS_INLINE friend bool operator!=(const Iterator& left, const Iterator& right) {
                if (const auto* left_index = std::get_if<size_t>(&left.it)) {
                    return *left_index != std::get<size_t>(right.it);
                }
                return std::get<DomIterator>(left.it) != std::get<DomIterator>(right.it);
            }
            ALWAYS_INLINE friend bool operator==(const Iterator& left, const Iterator& right) {
                return !(left != right);
            }

        private:
            std::variant<DomIterator, size_t> it;
            const std::vector<std::string>* keys = nullptr;
            const std::vector<Node>* values = nullptr;
        };
        ALWAYS_INLINE explicit Object(const simdjson::dom::object& object_)
                : dom_object(object_) {} /// NOLINT
        ALWAYS_INLINE explicit Object(const std::vector<std::string>* keys_,
                                      const std::vector<Node>* values_)
                : keys(keys_), values(values_) {} /// NOLINT
        ALWAYS_INLINE Iterator begin() const {
            return values ? Iterator(keys, values, 0) : Iterator(dom_object.begin());
        }
        ALWAYS_INLINE Iterator end() const {
            return values ? Iterator(keys, values, size()) : Iterator(dom_object.end());
        }
        ALWAYS_INLINE size_t size() const { return values ? values->size() : dom_object.size(); }
        /// Optional: Provides access to an object's element by index.
        KeyValuePair operator[](size_t index) const {
            assert(index < size());
            if (values) {
                return {(*keys)[index], Element(&(*values)[index])};
            }
            auto it = dom_object.begin();
            while (index--) {
                ++it;
            }
            const auto& res = *it;
            return {res.key, Element(res.value)};
        }

    private:
        simdjson::dom::object dom_object;
        const std::vector<std::string>* keys = nullptr;
        const std::vector<Node>* values = nullptr;
    };
    /// Parses a JSON document, returns the reference to its root element if succeeded.
    bool parse(const char* data, size_t size, Element& result, bool preserve_raw_numbers = false) {
        if (!preserve_raw_numbers) {
            auto document = dom_parser.parse(data, size);
            if (document.error()) {
                return false;
            }
            result = Element(document.value_unsafe());
            return true;
        }

        return parse_ondemand(data, size, result);
    }

private:
    bool parse_ondemand(const char* data, size_t size, Element& result) {
        simdjson::padded_string padded_json(data, size);
        simdjson::ondemand::document document;
        auto error = ondemand_parser.iterate(padded_json).get(document);
        if (error) {
            return false;
        }
        root = Node();
        if (!build_node(document, &root)) {
            return false;
        }
        result = Element(&root);
        return true;
    }

    static std::string_view trim_raw_number(std::string_view raw_number) {
        auto is_space = [](char ch) { return ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r'; };
        while (!raw_number.empty() && is_space(raw_number.front())) {
            raw_number.remove_prefix(1);
        }
        while (!raw_number.empty() && is_space(raw_number.back())) {
            raw_number.remove_suffix(1);
        }
        return raw_number;
    }

    template <typename RawNumber>
    static bool assign_raw_number(RawNumber&& raw_number, std::string* out) {
        if constexpr (std::is_same_v<std::decay_t<RawNumber>, std::string_view>) {
            *out = std::string(trim_raw_number(raw_number));
            return true;
        } else {
            std::string_view raw_number_view;
            auto error = std::move(raw_number).get(raw_number_view);
            if (error) {
                return false;
            }
            *out = std::string(trim_raw_number(raw_number_view));
            return true;
        }
    }

    template <typename Value>
    bool build_array_node(Value& value, Node* out) {
        simdjson::ondemand::array array;
        auto error = value.get_array().get(array);
        if (error) {
            return false;
        }
        out->type = Node::Type::ARRAY;
        for (auto element_result : array) {
            simdjson::ondemand::value element;
            error = std::move(element_result).get(element);
            if (error) {
                return false;
            }
            Node element_node;
            if (!build_node(element, &element_node)) {
                return false;
            }
            out->array_values.push_back(std::move(element_node));
        }
        return true;
    }

    template <typename Value>
    bool build_object_node(Value& value, Node* out) {
        simdjson::ondemand::object object;
        auto error = value.get_object().get(object);
        if (error) {
            return false;
        }
        out->type = Node::Type::OBJECT;
        for (auto field_result : object) {
            simdjson::ondemand::field field;
            error = std::move(field_result).get(field);
            if (error) {
                return false;
            }
            std::string_view key;
            error = field.unescaped_key().get(key);
            if (error) {
                return false;
            }
            std::string key_copy(key);
            simdjson::ondemand::value field_value = field.value();
            Node field_node;
            if (!build_node(field_value, &field_node)) {
                return false;
            }
            out->object_keys.push_back(std::move(key_copy));
            out->object_values.push_back(std::move(field_node));
        }
        return true;
    }

    template <typename Value>
    bool build_number_node(Value& value, Node* out) {
        simdjson::ondemand::number_type number_type;
        auto error = value.get_number_type().get(number_type);
        if (error) {
            return false;
        }
        switch (number_type) {
        case simdjson::ondemand::number_type::signed_integer:
            out->type = Node::Type::INT64;
            error = value.get_int64().get(out->int64_value);
            return !error;
        case simdjson::ondemand::number_type::unsigned_integer:
            out->type = Node::Type::UINT64;
            error = value.get_uint64().get(out->uint64_value);
            return !error;
        case simdjson::ondemand::number_type::floating_point_number:
            if (!assign_raw_number(value.raw_json_token(), &out->raw_number)) {
                return false;
            }
            out->type = Node::Type::DOUBLE;
            error = value.get_double().get(out->double_value);
            return !error;
        case simdjson::ondemand::number_type::big_integer: {
            if (!assign_raw_number(value.raw_json_token(), &out->raw_number)) {
                return false;
            }
            StringParser::ParseResult parse_result;
            out->int128_value = StringParser::string_to_int<Int128>(
                    out->raw_number.data(), out->raw_number.size(), &parse_result);
            out->type = parse_result == StringParser::PARSE_SUCCESS ? Node::Type::INT128
                                                                    : Node::Type::BIG_INTEGER;
            return true;
        }
        }
        return false;
    }

    template <typename Value>
    bool build_string_node(Value& value, Node* out) {
        std::string_view str;
        auto error = value.get_string().get(str);
        if (error) {
            return false;
        }
        out->type = Node::Type::STRING;
        out->string_value = std::string(str);
        return true;
    }

    template <typename Value>
    bool build_node(Value& value, Node* out) {
        simdjson::ondemand::json_type type;
        auto error = value.type().get(type);
        if (error) {
            return false;
        }
        switch (type) {
        case simdjson::ondemand::json_type::array:
            return build_array_node(value, out);
        case simdjson::ondemand::json_type::object:
            return build_object_node(value, out);
        case simdjson::ondemand::json_type::number:
            return build_number_node(value, out);
        case simdjson::ondemand::json_type::string: {
            return build_string_node(value, out);
        }
        case simdjson::ondemand::json_type::boolean:
            out->type = Node::Type::BOOL;
            error = value.get_bool().get(out->bool_value);
            return !error;
        case simdjson::ondemand::json_type::null:
            out->type = Node::Type::NULL_VALUE;
            return true;
        }
        return false;
    }

    simdjson::dom::parser dom_parser;
    simdjson::ondemand::parser ondemand_parser;
    Node root;
};
inline ALWAYS_INLINE SimdJSONParser::Array SimdJSONParser::Element::getArray() const {
    return node ? Array(&node->array_values) : Array(dom_element.get_array().value_unsafe());
}
inline ALWAYS_INLINE SimdJSONParser::Object SimdJSONParser::Element::getObject() const {
    return node ? Object(&node->object_keys, &node->object_values)
                : Object(dom_element.get_object().value_unsafe());
}

} // namespace doris

Coverage Report

Created: 2026-05-17 19:09