// 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.

#include <gtest/gtest.h>

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include "common/status.h"
#include "core/column/column_const.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_jsonb.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/data_type_string.h"
#include "core/data_type/define_primitive_type.h"
#include "core/data_type/primitive_type.h"
#include "core/data_type_serde/data_type_serde.h"
#include "core/types.h"
#include "exprs/function/function.h"
#include "exprs/function/function_test_util.h"
#include "exprs/function_context.h"
#include "testutil/any_type.h"
#include "testutil/column_helper.h"
#include "util/jsonb_writer.h"

namespace doris {
using namespace ut_type;

namespace {

ColumnPtr create_jsonb_column(size_t rows) {
    auto jsonb_type = std::make_shared<DataTypeJsonb>();
    auto jsonb_column = jsonb_type->create_column();
    for (size_t i = 0; i < rows; ++i) {
        EXPECT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":{"b":1}})")));
    }
    return jsonb_column;
}

ColumnPtr create_path_column(const std::vector<std::string>& paths) {
    auto path_type = std::make_shared<DataTypeString>();
    auto path_column = path_type->create_column();
    for (const auto& path : paths) {
        EXPECT_TRUE(insert_cell(path_column, path_type, STRING(path)));
    }
    return path_column;
}

ColumnPtr create_const_path_column(const std::string& path, size_t rows) {
    return ColumnConst::create(create_path_column({path}), rows);
}

ColumnPtr create_const_nullable_path_column(const std::string& path, size_t rows) {
    auto nullable_path = ColumnNullable::create(ColumnString::create(), ColumnUInt8::create());
    nullable_path->insert_data(path.data(), path.size());
    return ColumnConst::create(std::move(nullable_path), rows);
}

Status execute_json_keys_with_path(ColumnPtr path_column, size_t rows) {
    auto jsonb_type = std::make_shared<DataTypeJsonb>();
    auto path_type = std::make_shared<DataTypeString>();
    auto return_type = make_nullable(
            std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));

    Block block;
    block.insert({create_jsonb_column(rows), jsonb_type, "jsonb"});
    block.insert({std::move(path_column), path_type, "path"});

    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            "json_keys", block.get_columns_with_type_and_name(), return_type);
    DORIS_CHECK(func != nullptr);

    FunctionUtils fn_utils(return_type, {jsonb_type, path_type}, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    RETURN_IF_ERROR(func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
    RETURN_IF_ERROR(func->open(fn_ctx, FunctionContext::THREAD_LOCAL));

    block.insert({nullptr, return_type, "result"});
    auto st = func->execute(fn_ctx, block, {0, 1}, block.columns() - 1, rows);

    static_cast<void>(func->close(fn_ctx, FunctionContext::THREAD_LOCAL));
    static_cast<void>(func->close(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
    return st;
}

void expect_invalid_json_keys_super_wildcard_path(const Status& st) {
    EXPECT_EQ(st.code(), ErrorCode::INVALID_JSON_PATH) << st.to_string();
    EXPECT_NE(st.to_string().find(
                      "path expressions may not contain the * and ** tokens or an array range"),
              std::string::npos)
            << st.to_string();
}

} // namespace

TEST(FunctionJsonbTEST, JsonbParseTest) {
    std::string func_name = "json_parse";
    InputTypeSet input_types = {Nullable {PrimitiveType::TYPE_VARCHAR}};

    DataSet data_set_valid = {
            {{Null()}, Null()},
            {{STRING("null")}, STRING("null")},
            {{STRING("true")}, STRING("true")},
            {{STRING("false")}, STRING("false")},
            {{STRING("100")}, STRING("100")},                                 //int8
            {{STRING("10000")}, STRING("10000")},                             // int16
            {{STRING("1000000000")}, STRING("1000000000")},                   // int32
            {{STRING("1152921504606846976")}, STRING("1152921504606846976")}, // int64
            {{STRING("6.18")}, STRING("6.18")},                               // double
            {{STRING(R"("abcd")")}, STRING(R"("abcd")")},                     // string
            {{STRING("{}")}, STRING("{}")},                                   // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})")}, STRING(R"({"k1":"v31","k2":300})")}, // object
            {{STRING("[]")}, STRING("[]")},                              // empty array
            {{STRING("[123, 456]")}, STRING("[123,456]")},               // int array
            {{STRING(R"(["abc", "def"])")}, STRING(R"(["abc","def"])")}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])")},
             STRING(R"([null,true,false,100,6.18,"abc"])")}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])")},
             STRING(R"([{"k1":"v41","k2":400},1,"a",3.14])")}, // complex array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_valid));

    DataSet data_set_invalid = {
            {{STRING("abc")}, Null()}, // invalid string
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));

    data_set_invalid = {
            {{STRING("'abc'")}, Null()}, // invalid string
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));

    data_set_invalid = {
            {{STRING("100x")}, Null()}, // invalid int
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));

    data_set_invalid = {
            {{STRING("6.a8")}, Null()}, // invalid double
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));

    data_set_invalid = {
            {{STRING("{x")}, Null()}, // invalid object
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));

    data_set_invalid = {
            {{STRING("[123, abc]")}, Null()} // invalid array
    };
    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set_invalid,
                                                          -1, -1, true));
}

TEST(FunctionJsonbTEST, JsonbParseErrorToNullTest) {
    std::string func_name = "json_parse_error_to_null";
    InputTypeSet input_types = {Nullable {PrimitiveType::TYPE_VARCHAR}};

    DataSet data_set = {
            {{Null()}, Null()},
            {{STRING("null")}, STRING("null")},
            {{STRING("true")}, STRING("true")},
            {{STRING("false")}, STRING("false")},
            {{STRING("100")}, STRING("100")},                                 //int8
            {{STRING("10000")}, STRING("10000")},                             // int16
            {{STRING("1000000000")}, STRING("1000000000")},                   // int32
            {{STRING("1152921504606846976")}, STRING("1152921504606846976")}, // int64
            {{STRING("6.18")}, STRING("6.18")},                               // double
            {{STRING(R"("abcd")")}, STRING(R"("abcd")")},                     // string
            {{STRING("{}")}, STRING("{}")},                                   // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})")}, STRING(R"({"k1":"v31","k2":300})")}, // object
            {{STRING("[]")}, STRING("[]")},                              // empty array
            {{STRING("[123, 456]")}, STRING("[123,456]")},               // int array
            {{STRING(R"(["abc", "def"])")}, STRING(R"(["abc","def"])")}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])")},
             STRING(R"([null,true,false,100,6.18,"abc"])")}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])")},
             STRING(R"([{"k1":"v41","k2":400},1,"a",3.14])")}, // complex array
            {{STRING("abc")}, Null()},                         // invalid string
            {{STRING("'abc'")}, Null()},                       // invalid string
            {{STRING("100x")}, Null()},                        // invalid int
            {{STRING("6.a8")}, Null()},                        // invalid double
            {{STRING("{x")}, Null()},                          // invalid object
            {{STRING("[123, abc]")}, Null()}                   // invalid array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));
}

TEST(FunctionJsonbTEST, JsonbParseErrorToValueTest) {
    std::string func_name = "json_parse_error_to_value";
    InputTypeSet input_types = {Nullable {PrimitiveType::TYPE_VARCHAR}, PrimitiveType::TYPE_JSONB};

    DataSet data_set = {
            {{Null(), STRING("{}")}, Null()},
            {{STRING("null"), STRING("{}")}, STRING("null")},
            {{STRING("true"), STRING("{}")}, STRING("true")},
            {{STRING("false"), STRING("{}")}, STRING("false")},
            {{STRING("100"), STRING("{}")}, STRING("100")},                                 //int8
            {{STRING("10000"), STRING("{}")}, STRING("10000")},                             // int16
            {{STRING("1000000000"), STRING("{}")}, STRING("1000000000")},                   // int32
            {{STRING("1152921504606846976"), STRING("{}")}, STRING("1152921504606846976")}, // int64
            {{STRING("6.18"), STRING("{}")}, STRING("6.18")},           // double
            {{STRING(R"("abcd")"), STRING("{}")}, STRING(R"("abcd")")}, // string
            {{STRING("{}"), STRING("{}")}, STRING("{}")},               // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("{}")},
             STRING(R"({"k1":"v31","k2":300})")},                        // object
            {{STRING("[]"), STRING("{}")}, STRING("[]")},                // empty array
            {{STRING("[123, 456]"), STRING("{}")}, STRING("[123,456]")}, // int array
            {{STRING(R"(["abc", "def"])"), STRING("{}")},
             STRING(R"(["abc","def"])")}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("{}")},
             STRING(R"([null,true,false,100,6.18,"abc"])")}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("{}")},
             STRING(R"([{"k1":"v41","k2":400},1,"a",3.14])")},           // complex array
            {{STRING("abc"), STRING(R"("abc")")}, STRING(R"("abc")")},   // invalid string
            {{STRING("'abc'"), STRING(R"("abc")")}, STRING(R"("abc")")}, // invalid string
            {{STRING("100x"), STRING("100")}, STRING("100")},            // invalid int
            {{STRING("6.a8"), STRING("6.18")}, STRING("6.18")},          // invalid double
            {{STRING("{x"), STRING("{}")}, STRING("{}")},                // invalid object
            {{STRING("[123, abc]"), STRING(R"([123,"abc"])")},
             STRING(R"([123,"abc"])")} // invalid array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));

    InputTypeSet input_types2 = {Nullable {PrimitiveType::TYPE_VARCHAR}, PrimitiveType::TYPE_JSONB,
                                 PrimitiveType::TYPE_JSONB};
    DataSet data_set2 = {{{Null(), STRING("{}"), STRING("{}")}, Null()},
                         {{STRING("{}"), STRING("{}"), STRING("{}")}, Null()}};
    auto st = check_function<DataTypeJsonb, true>(func_name, input_types2, data_set2);
    ASSERT_EQ(st.code(), ErrorCode::INVALID_ARGUMENT) << st.to_string();
}

TEST(FunctionJsonbTEST, JsonbKeysRejectSuperWildcardPath) {
    auto const_path_status = execute_json_keys_with_path(create_const_path_column("$**.a", 1), 1);
    expect_invalid_json_keys_super_wildcard_path(const_path_status);

    auto non_const_path_status =
            execute_json_keys_with_path(create_path_column({"$.a", "$**.a"}), 2);
    expect_invalid_json_keys_super_wildcard_path(non_const_path_status);
}

TEST(FunctionJsonbTEST, JsonbExtractTest) {
    std::string func_name = "jsonb_extract";
    InputTypeSet input_types = {PrimitiveType::TYPE_JSONB, PrimitiveType::TYPE_VARCHAR};

    // json_extract root
    DataSet data_set = {
            {{Null(), STRING("$")}, Null()},
            {{STRING("null"), STRING("$")}, STRING("null")},
            {{STRING("true"), STRING("$")}, STRING("true")},
            {{STRING("false"), STRING("$")}, STRING("false")},
            {{STRING("100"), STRING("$")}, STRING("100")},                                 //int8
            {{STRING("10000"), STRING("$")}, STRING("10000")},                             // int16
            {{STRING("1000000000"), STRING("$")}, STRING("1000000000")},                   // int32
            {{STRING("1152921504606846976"), STRING("$")}, STRING("1152921504606846976")}, // int64
            {{STRING("6.18"), STRING("$")}, STRING("6.18")},                               // double
            {{STRING(R"("abcd")"), STRING("$")}, STRING(R"("abcd")")},                     // string
            {{STRING("{}"), STRING("$")}, STRING("{}")}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("$")},
             STRING(R"({"k1":"v31","k2":300})")},                       // object
            {{STRING("[]"), STRING("$")}, STRING("[]")},                // empty array
            {{STRING("[123, 456]"), STRING("$")}, STRING("[123,456]")}, // int array
            {{STRING(R"(["abc", "def"])"), STRING("$")},
             STRING(R"(["abc","def"])")}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$")},
             STRING(R"([null,true,false,100,6.18,"abc"])")}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$")},
             STRING(R"([{"k1":"v41","k2":400},1,"a",3.14])")}, // complex array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));

    // json_extract obejct
    data_set = {
            {{Null(), STRING("$.k1")}, Null()},
            {{STRING("null"), STRING("$.k1")}, Null()},
            {{STRING("true"), STRING("$.k1")}, Null()},
            {{STRING("false"), STRING("$.k1")}, Null()},
            {{STRING("100"), STRING("$.k1")}, Null()},                 //int8
            {{STRING("10000"), STRING("$.k1")}, Null()},               // int16
            {{STRING("1000000000"), STRING("$.k1")}, Null()},          // int32
            {{STRING("1152921504606846976"), STRING("$.k1")}, Null()}, // int64
            {{STRING("6.18"), STRING("$.k1")}, Null()},                // double
            {{STRING(R"("abcd")"), STRING("$.k1")}, Null()},           // string
            {{STRING("{}"), STRING("$.k1")}, Null()},                  // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("$.k1")}, STRING(R"("v31")")}, // object
            {{STRING("[]"), STRING("$.k1")}, Null()},                // empty array
            {{STRING("[123, 456]"), STRING("$.k1")}, Null()},        // int array
            {{STRING(R"(["abc", "def"])"), STRING("$.k1")}, Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$.k1")},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$.k1")},
             Null()}, // complex array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));

    // json_extract array
    data_set = {
            {{Null(), STRING("$[0]")}, Null()},
            {{STRING("null"), STRING("$[0]")}, STRING("null")},
            {{STRING("true"), STRING("$[0]")}, STRING("true")},
            {{STRING("false"), STRING("$[0]")}, STRING("false")},
            {{STRING("100"), STRING("$[0]")}, STRING("100")},               //int8
            {{STRING("10000"), STRING("$[0]")}, STRING("10000")},           // int16
            {{STRING("1000000000"), STRING("$[0]")}, STRING("1000000000")}, // int32
            {{STRING("1152921504606846976"), STRING("$[0]")},
             STRING("1152921504606846976")},                             // int64
            {{STRING("6.18"), STRING("$[0]")}, STRING("6.18")},          // double
            {{STRING(R"("abcd")"), STRING("$[0]")}, STRING("\"abcd\"")}, // string
            {{STRING("{}"), STRING("$[0]")}, STRING("{}")},              // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("$[0]")},
             STRING(R"({"k1":"v31","k2":300})")},     // object
            {{STRING("[]"), STRING("$[0]")}, Null()}, // empty array
            {{STRING("null"), STRING("$[1]")}, Null()},
            {{STRING("true"), STRING("$[1]")}, Null()},
            {{STRING("false"), STRING("$[1]")}, Null()},
            {{STRING("100"), STRING("$[1]")}, Null()},                           //int8
            {{STRING("10000"), STRING("$[1]")}, Null()},                         // int16
            {{STRING("1000000000"), STRING("$[1]")}, Null()},                    // int32
            {{STRING("1152921504606846976"), STRING("$[1]")}, Null()},           // int64
            {{STRING("6.18"), STRING("$[1]")}, Null()},                          // double
            {{STRING(R"("abcd")"), STRING("$[1]")}, Null()},                     // string
            {{STRING("{}"), STRING("$[1]")}, Null()},                            // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("$[1]")}, Null()},    // object
            {{STRING("[]"), STRING("$[1]")}, Null()},                            // empty array
            {{STRING("[123, 456]"), STRING("$[0]")}, STRING("123")},             // int array
            {{STRING("[123, 456]"), STRING("$[1]")}, STRING("456")},             // int array
            {{STRING("[123, 456]"), STRING("$[2]")}, Null()},                    // int array
            {{STRING(R"(["abc", "def"])"), STRING("$[0]")}, STRING(R"("abc")")}, // string array
            {{STRING(R"(["abc", "def"])"), STRING("$[1]")}, STRING(R"("def")")}, // string array
            {{STRING(R"(["abc", "def"])"), STRING("$[2]")}, Null()},             // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[0]")},
             STRING("null")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[1]")},
             STRING("true")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[2]")},
             STRING("false")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[3]")},
             STRING("100")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[4]")},
             STRING("6.18")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[5]")},
             STRING(R"("abc")")}, // multi type array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[6]")},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[0]")},
             STRING(R"({"k1":"v41","k2":400})")}, // complex array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[1]")},
             STRING("1")}, // complex array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[2]")},
             STRING(R"("a")")}, // complex array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[3]")},
             STRING("3.14")}, // complex array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[4]")},
             Null()}, // complex array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));

    // json_extract $[0].k1
    data_set = {
            {{Null(), STRING("$[0].k1")}, Null()},
            {{STRING("null"), STRING("$[0].k1")}, Null()},
            {{STRING("true"), STRING("$[0].k1")}, Null()},
            {{STRING("false"), STRING("$[0].k1")}, Null()},
            {{STRING("100"), STRING("$[0].k1")}, Null()},                 //int8
            {{STRING("10000"), STRING("$[0].k1")}, Null()},               // int16
            {{STRING("1000000000"), STRING("$[0].k1")}, Null()},          // int32
            {{STRING("1152921504606846976"), STRING("$[0].k1")}, Null()}, // int64
            {{STRING("6.18"), STRING("$[0].k1")}, Null()},                // double
            {{STRING(R"("abcd")"), STRING("$[0].k1")}, Null()},           // string
            {{STRING("{}"), STRING("$[0].k1")}, Null()},                  // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("$[0].k1")},
             STRING(R"("v31")")},                                       // object
            {{STRING("[]"), STRING("$[0].k1")}, Null()},                // empty array
            {{STRING("[123, 456]"), STRING("$[0].k1")}, Null()},        // int array
            {{STRING(R"(["abc", "def"])"), STRING("$[0].k1")}, Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("$[0].k1")},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("$[0].k1")},
             STRING(R"("v41")")}, // complex array
    };

    static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));
}

TEST(FunctionJsonbTEST, JsonExtractCheckArg) {
    ColumnsWithTypeAndName args;
    args.emplace_back(ColumnWithTypeAndName {nullptr, std::make_shared<DataTypeJsonb>(), "jsonb"});
    args.emplace_back(ColumnWithTypeAndName {nullptr, std::make_shared<DataTypeString>(), "path"});
    auto return_type = make_nullable(std::make_shared<DataTypeJsonb>());
    FunctionBasePtr func =
            SimpleFunctionFactory::instance().get_function("json_extract", args, return_type);
    ASSERT_NE(func, nullptr);

    auto function_ptr = std::dynamic_pointer_cast<IFunction>(
            assert_cast<DefaultFunction*>(func.get())->function);
    ASSERT_TRUE(function_ptr);

    ASSERT_EQ(function_ptr->get_number_of_arguments(), 0);
    ASSERT_TRUE(function_ptr->is_variadic());

    FunctionUtils fn_utils(return_type, {}, false);

    auto st = function_ptr->open(fn_utils.get_fn_ctx(),
                                 FunctionContext::FunctionStateScope::THREAD_LOCAL);
    ASSERT_TRUE(st.ok()) << st.to_string();

    auto col1 = ColumnHelper::create_column<DataTypeString>({"ab", "cd"});
    auto col2 = ColumnHelper::create_column<DataTypeString>({"ef", "gh"});

    Block block;
    block.insert({col1, std::make_shared<DataTypeString>(), "jsonb"});
    block.insert({col2, std::make_shared<DataTypeString>(), "path"});

    st = function_ptr->execute(fn_utils.get_fn_ctx(), block, {0, 1}, 2, 2);
    // the 1st arg's primitive type should be jsonb, but here is string.
    ASSERT_EQ(st.code(), ErrorCode::INVALID_ARGUMENT);
}

TEST(FunctionJsonbTEST, JsonParseCheckArg) {
    ColumnsWithTypeAndName args;
    args.emplace_back(
            ColumnWithTypeAndName {nullptr, std::make_shared<DataTypeString>(), "json_str"});
    args.emplace_back(
            ColumnWithTypeAndName {nullptr, std::make_shared<DataTypeJsonb>(), "default_value"});
    auto return_type = std::make_shared<DataTypeJsonb>();
    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            "json_parse_error_to_value", args, return_type);
    ASSERT_NE(func, nullptr);

    auto function_ptr = std::dynamic_pointer_cast<IFunction>(
            assert_cast<DefaultFunction*>(func.get())->function);
    ASSERT_TRUE(function_ptr);

    ASSERT_EQ(function_ptr->get_number_of_arguments(), 0);
    ASSERT_TRUE(function_ptr->is_variadic());

    FunctionUtils fn_utils(return_type,
                           {std::make_shared<DataTypeString>(), std::make_shared<DataTypeJsonb>()},
                           false);

    auto st = function_ptr->open(fn_utils.get_fn_ctx(),
                                 FunctionContext::FunctionStateScope::THREAD_LOCAL);
    ASSERT_TRUE(st.ok()) << st.to_string();

    auto col1 = ColumnHelper::create_column<DataTypeString>({"ab", "cd"});
    auto col2 = ColumnHelper::create_column<DataTypeString>({"ef", "gh"});

    Block block;
    block.insert({col1, std::make_shared<DataTypeString>(), "json_str"});
    block.insert({col2, std::make_shared<DataTypeString>(), "default_value"});

    st = function_ptr->execute(fn_utils.get_fn_ctx(), block, {0, 1}, 2, 2);
    // the 2nd arg's primitive type should be jsonb, but here is string.
    ASSERT_EQ(st.code(), ErrorCode::INVALID_ARGUMENT);
}

TEST(FunctionJsonbTEST, JsonbCastToOtherTest) {
    std::string func_name = "CAST";
    InputTypeSet input_types = {Nullable {PrimitiveType::TYPE_JSONB},
                                Consted {PrimitiveType::TYPE_BOOLEAN}};

    // cast to boolean
    DataSet data_set = {
            {{STRING("null"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)}, Null()},
            {{STRING("true"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)}, BOOLEAN(1)},
            {{STRING("false"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)}, BOOLEAN(0)},
            {{STRING("100"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)}, BOOLEAN(1)}, //int8
            {{STRING("10000"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             BOOLEAN(1)}, // int16
            {{STRING("1000000000"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             BOOLEAN(1)}, // int32
            {{STRING("1152921504606846976"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             BOOLEAN(1)}, // int64
            {{STRING("6.18"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             BOOLEAN(1)}, // double
            {{STRING(R"("abcd")"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // string
            {{STRING("{}"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"),
              static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // object
            {{STRING("[]"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // empty array
            {{STRING("[123, 456]"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // int array
            {{STRING(R"(["abc", "def"])"), static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"),
              static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"),
              static_cast<uint8_t>(PrimitiveType::TYPE_BOOLEAN)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeUInt8, true>(func_name, input_types, const_dataset));
    }
    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_TINYINT}};
    // cast to TINYINT
    data_set = {
            {{STRING("null"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)}, Null()},
            {{STRING("true"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)}, TINYINT(1)},
            {{STRING("false"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)}, TINYINT(0)},
            {{STRING("100"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             TINYINT(100)},                                                                //int8
            {{STRING("10000"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)}, Null()}, // int16
            {{STRING("1000000000"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // int32
            {{STRING("1152921504606846976"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // int64
            {{STRING("6.18"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             TINYINT(6)}, // double
            {{STRING(R"("abcd")"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // string
            {{STRING("{}"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"),
              static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // object
            {{STRING("[]"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // empty array
            {{STRING("[123, 456]"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // int array
            {{STRING(R"(["abc", "def"])"), static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"),
              static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"),
              static_cast<int8_t>(PrimitiveType::TYPE_TINYINT)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeInt8, true>(func_name, input_types, const_dataset));
    }

    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_SMALLINT}};
    // cast to SMALLINT
    data_set = {
            {{STRING("null"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)}, Null()},
            {{STRING("true"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)}, SMALLINT(1)},
            {{STRING("false"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)}, SMALLINT(0)},
            {{STRING("100"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             SMALLINT(100)}, //int8
            {{STRING("10000"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             SMALLINT(10000)}, // int16
            {{STRING("1000000000"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // int32
            {{STRING("1152921504606846976"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // int64
            {{STRING("6.18"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             SMALLINT(6)}, // double
            {{STRING(R"("abcd")"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // string
            {{STRING("{}"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"),
              static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // object
            {{STRING("[]"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // empty array
            {{STRING("[123, 456]"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // int array
            {{STRING(R"(["abc", "def"])"), static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"),
              static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"),
              static_cast<int16_t>(PrimitiveType::TYPE_SMALLINT)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeInt16, true>(func_name, input_types, const_dataset));
    }

    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_INT}};
    // cast to INT
    data_set = {
            {{STRING("null"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, Null()},
            {{STRING("true"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, INT(1)},
            {{STRING("false"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, INT(0)},
            {{STRING("100"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, INT(100)},     //int8
            {{STRING("10000"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, INT(10000)}, // int16
            {{STRING("1000000000"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             INT(1000000000)}, // int32
            {{STRING("1152921504606846976"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()},                                                                  // int64
            {{STRING("6.18"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, INT(6)}, // double
            {{STRING(R"("abcd")"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()},                                                                // string
            {{STRING("{}"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, Null()}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()},                                                                // object
            {{STRING("[]"), static_cast<int32_t>(PrimitiveType::TYPE_INT)}, Null()}, // empty array
            {{STRING("[123, 456]"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()}, // int array
            {{STRING(R"(["abc", "def"])"), static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"),
              static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"),
              static_cast<int32_t>(PrimitiveType::TYPE_INT)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeInt32, true>(func_name, input_types, const_dataset));
    }

    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_BIGINT}};
    // cast to BIGINT
    data_set = {
            {{STRING("null"), BIGINT(1)}, Null()},
            {{STRING("true"), BIGINT(1)}, BIGINT(1)},
            {{STRING("false"), BIGINT(1)}, BIGINT(0)},
            {{STRING("100"), BIGINT(1)}, BIGINT(100)},                                 //int8
            {{STRING("10000"), BIGINT(1)}, BIGINT(10000)},                             // int16
            {{STRING("1000000000"), BIGINT(1)}, BIGINT(1000000000)},                   // int32
            {{STRING("1152921504606846976"), BIGINT(1)}, BIGINT(1152921504606846976)}, // int64
            {{STRING("6.18"), BIGINT(1)}, BIGINT(6)},                                  // double
            {{STRING(R"("abcd")"), BIGINT(1)}, Null()},                                // string
            {{STRING("{}"), BIGINT(1)}, Null()},                         // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), BIGINT(1)}, Null()}, // object
            {{STRING("[]"), BIGINT(1)}, Null()},                         // empty array
            {{STRING("[123, 456]"), BIGINT(1)}, Null()},                 // int array
            {{STRING(R"(["abc", "def"])"), BIGINT(1)}, Null()},          // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), BIGINT(1)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), BIGINT(1)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeInt64, true>(func_name, input_types, const_dataset));
    }

    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_DOUBLE}};
    // cast to DOUBLE
    data_set = {
            {{STRING("null"), DOUBLE(1)}, Null()},
            {{STRING("true"), DOUBLE(1)}, DOUBLE(1)},
            {{STRING("false"), DOUBLE(1)}, DOUBLE(0)},
            {{STRING("100"), DOUBLE(1)}, DOUBLE(100)},                                 //int8
            {{STRING("10000"), DOUBLE(1)}, DOUBLE(10000)},                             // int16
            {{STRING("1000000000"), DOUBLE(1)}, DOUBLE(1000000000)},                   // int32
            {{STRING("1152921504606846976"), DOUBLE(1)}, DOUBLE(1152921504606846976)}, // int64
            {{STRING("6.18"), DOUBLE(1)}, DOUBLE(6.18)},                               // double
            {{STRING(R"("abcd")"), DOUBLE(1)}, Null()},                                // string
            {{STRING("{}"), DOUBLE(1)}, Null()},                         // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), DOUBLE(1)}, Null()}, // object
            {{STRING("[]"), DOUBLE(1)}, Null()},                         // empty array
            {{STRING("[123, 456]"), DOUBLE(1)}, Null()},                 // int array
            {{STRING(R"(["abc", "def"])"), DOUBLE(1)}, Null()},          // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), DOUBLE(1)},
             Null()}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), DOUBLE(1)},
             Null()}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeFloat64, true>(func_name, input_types, const_dataset));
    }

    input_types = {Nullable {PrimitiveType::TYPE_JSONB}, Consted {PrimitiveType::TYPE_VARCHAR}};
    // cast to STRING
    data_set = {
            {{STRING("null"), STRING("1")}, Null()},
            {{STRING("true"), STRING("1")}, STRING("true")},
            {{STRING("false"), STRING("1")}, STRING("false")},
            {{STRING("100"), STRING("1")}, STRING("100")},                                 //int8
            {{STRING("10000"), STRING("1")}, STRING("10000")},                             // int16
            {{STRING("1000000000"), STRING("1")}, STRING("1000000000")},                   // int32
            {{STRING("1152921504606846976"), STRING("1")}, STRING("1152921504606846976")}, // int64
            {{STRING("6.18"), STRING("1")}, STRING("6.18")},                               // double
            {{STRING(R"("abcd")"), STRING("1")}, STRING(R"(abcd)")},                       // string
            {{STRING("{}"), STRING("1")}, STRING("{}")}, // empty object
            {{STRING(R"({"k1":"v31", "k2": 300})"), STRING("1")},
             STRING(R"({"k1":"v31","k2":300})")},                       // object
            {{STRING("[]"), STRING("1")}, STRING("[]")},                // empty array
            {{STRING("[123, 456]"), STRING("1")}, STRING("[123,456]")}, // int array
            {{STRING(R"(["abc", "def"])"), STRING("1")},
             STRING(R"(["abc","def"])")}, // string array
            {{STRING(R"([null, true, false, 100, 6.18, "abc"])"), STRING("1")},
             STRING(R"([null,true,false,100,6.18,"abc"])")}, // multi type array
            {{STRING(R"([{"k1":"v41", "k2": 400}, 1, "a", 3.14])"), STRING("1")},
             STRING(R"([{"k1":"v41","k2":400},1,"a",3.14])")}, // complex array
    };
    for (const auto& row : data_set) {
        DataSet const_dataset = {row};
        static_cast<void>(
                check_function<DataTypeString, true>(func_name, input_types, const_dataset));
    }
}

TEST(FunctionJsonbTEST, JsonbCastFromOtherTest) {
    // CAST Nullable(X) to Nullable(JSONB)
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_BOOLEAN}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{BOOLEAN(1), Null()}, STRING("true")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_BOOLEAN}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{BOOLEAN(0), Null()}, STRING("false")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_TINYINT}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{TINYINT(100), Null()}, STRING("100")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_SMALLINT}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{SMALLINT(10000), Null()}, STRING("10000")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_INT}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{INT(1000000000), Null()}, STRING("1000000000")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_BIGINT}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{BIGINT(1152921504606846976), Null()}, STRING("1152921504606846976")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_DOUBLE}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{DOUBLE(6.18), Null()}, STRING("6.18")}}));
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_VARCHAR}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{STRING(R"(abcd)"), Null()}, Null()}})); // should fail
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST", {Nullable {PrimitiveType::TYPE_VARCHAR}, Consted {PrimitiveType::TYPE_JSONB}},
            {{{STRING(R"("abcd")"), Null()}, STRING(R"("abcd")")}}));

    // CAST X to JSONB. the second argument is just a dummy because result type is not nullable so we need it
    // rather than a Null.
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_BOOLEAN}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{BOOLEAN(1), STRING()}, STRING("true")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_BOOLEAN}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{BOOLEAN(0), STRING()}, STRING("false")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_TINYINT}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{TINYINT(100), STRING()}, STRING("100")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_SMALLINT}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{SMALLINT(10000), STRING()}, STRING("10000")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST", {Notnull {PrimitiveType::TYPE_INT}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{INT(1000000000), STRING()}, STRING("1000000000")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_BIGINT}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{BIGINT(1152921504606846976), STRING()}, STRING("1152921504606846976")}}));
    static_cast<void>(check_function<DataTypeJsonb, false>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_DOUBLE}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{DOUBLE(6.18), STRING()}, STRING("6.18")}}));
    // String to JSONB should always be Nullable
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_VARCHAR}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{STRING(R"(abcd)"), STRING()}, Null()}})); // should fail
    static_cast<void>(check_function<DataTypeJsonb, true>(
            "CAST",
            {Notnull {PrimitiveType::TYPE_VARCHAR}, ConstedNotnull {PrimitiveType::TYPE_JSONB}},
            {{{STRING(R"("abcd")"), STRING()}, STRING(R"("abcd")")}}));
}

TEST(FunctionJsonbTEST, JsonbToJson) {
    std::string func_name = "to_json";

    {
        InputTypeSet input_types = {PrimitiveType::TYPE_BOOLEAN};

        DataSet data_set = {
                {{Null()}, Null()},
                {{UInt8(1)}, STRING("true")},
                {{UInt8(0)}, STRING("false")},
        };
        static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, data_set));
    }

    static_cast<void>(
            check_function<DataTypeJsonb, true>("to_json", {Nullable {PrimitiveType::TYPE_INT}},
                                                {{{INT(1000000000)}, STRING("1000000000")}}));

    static_cast<void>(
            check_function<DataTypeJsonb, true>("to_json", {Nullable {PrimitiveType::TYPE_VARCHAR}},
                                                {{{STRING("hello")}, STRING(R"("hello")")}}));
}

TEST(FunctionJsonbTEST, JsonArray) {
    std::string func_name = "json_array";

    InputTypeSet input_types = {PrimitiveType::TYPE_JSONB};

    DataSet data_set = {
            {{Null()}, STRING("[null]")},
            {{STRING("null")}, STRING("[null]")},
            {{STRING("true")}, STRING("[true]")},
            {{STRING("false")}, STRING("[false]")},
            {{STRING("100")}, STRING("[100]")},                                 //int8
            {{STRING("10000")}, STRING("[10000]")},                             // int16
            {{STRING("1000000000")}, STRING("[1000000000]")},                   // int32
            {{STRING("1152921504606846976")}, STRING("[1152921504606846976]")}, // int64
            {{STRING("6.18")}, STRING("[6.18]")},                               // double
            {{STRING(R"("abcd")")}, STRING(R"(["abcd"])")},                     // string
    };

    static_cast<void>(check_function<DataTypeJsonb>(func_name, input_types, data_set));

    func_name = "json_array_ignore_null";

    data_set = {
            {{Null()}, STRING("[]")},
            {{STRING("null")}, STRING("[null]")},
            {{STRING("true")}, STRING("[true]")},
            {{STRING("false")}, STRING("[false]")},
            {{STRING("100")}, STRING("[100]")},                                 //int8
            {{STRING("10000")}, STRING("[10000]")},                             // int16
            {{STRING("1000000000")}, STRING("[1000000000]")},                   // int32
            {{STRING("1152921504606846976")}, STRING("[1152921504606846976]")}, // int64
            {{STRING("6.18")}, STRING("[6.18]")},                               // double
            {{STRING(R"("abcd")")}, STRING(R"(["abcd"])")},                     // string
    };

    static_cast<void>(check_function<DataTypeJsonb>(func_name, input_types, data_set));
}

TEST(FunctionJsonbTEST, JsonLength) {
    std::string func_name = "json_length";

    {
        InputTypeSet input_types = {PrimitiveType::TYPE_JSONB, PrimitiveType::TYPE_VARCHAR};

        DataSet data_set = {
                {{STRING(R"({"k1":"v1", "k2": 2})"), Null()}, Null()},
                {{STRING(R"([1, 2, 3])"), STRING("$")}, INT(3)},
                {{STRING(R"("string")"), STRING("$")}, INT(1)},
                {{STRING("null"), STRING("$")}, INT(1)},
        };
        static_cast<void>(check_function<DataTypeInt32, true>(func_name, input_types, data_set));
    }

    auto json_data_type = std::make_shared<DataTypeJsonb>();
    auto json_column = json_data_type->create_column();

    JsonbWriter writer;
    writer.writeStartArray();
    writer.writeString("hello");
    writer.writeString("world");
    writer.writeEndArray();

    json_column->insert_data(writer.getOutput()->getBuffer(), writer.getOutput()->getSize());

    Block block;
    block.insert(ColumnWithTypeAndName(std::move(json_column), json_data_type, "json_col"));

    auto path_column_ = ColumnNullable::create(ColumnString::create(), ColumnUInt8::create());
    path_column_->insert_default();
    auto path_column = ColumnConst::create(std::move(path_column_), 1);
    block.insert(ColumnWithTypeAndName(
            std::move(path_column),
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeString>()), "path_col"));

    auto return_type = make_nullable(std::make_shared<DataTypeInt32>());
    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            func_name, block.get_columns_with_type_and_name(), return_type);
    ASSERT_TRUE(func != nullptr);

    std::vector<DataTypePtr> arg_types;
    arg_types.emplace_back(json_data_type);
    arg_types.emplace_back(std::make_shared<DataTypeNullable>(std::make_shared<DataTypeString>()));

    FunctionUtils fn_utils(return_type, arg_types, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    //     fn_ctx->set_constant_cols(constant_cols);
    auto st = func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL);
    ASSERT_TRUE(st.ok()) << "open failed: " << st.to_string();
    st = func->open(fn_ctx, FunctionContext::THREAD_LOCAL);
    ASSERT_TRUE(st.ok()) << "open failed: " << st.to_string();

    block.insert({nullptr, return_type, "result"});

    auto result = block.columns() - 1;
    st = func->execute(fn_ctx, block, {0, 1}, result, 1);
    ASSERT_TRUE(st.ok()) << "execute failed: " << st.to_string();
}

TEST(FunctionJsonbTEST, JsonLengthConstNullablePathMultiRows) {
    constexpr size_t rows = 3;
    auto jsonb_type = std::make_shared<DataTypeJsonb>();
    auto path_type = make_nullable(std::make_shared<DataTypeString>());
    auto return_type = make_nullable(std::make_shared<DataTypeInt32>());

    auto jsonb_column = jsonb_type->create_column();
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[1,2]})")));
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[3,4,5]})")));
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[]})")));

    Block block;
    block.insert({std::move(jsonb_column), jsonb_type, "jsonb"});
    block.insert({create_const_nullable_path_column("$.a", rows), path_type, "path"});

    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            "json_length", block.get_columns_with_type_and_name(), return_type);
    ASSERT_NE(func, nullptr);

    FunctionUtils fn_utils(return_type, {jsonb_type, path_type}, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL).ok());
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::THREAD_LOCAL).ok());

    block.insert({nullptr, return_type, "result"});
    ASSERT_TRUE(func->execute(fn_ctx, block, {0, 1}, 2, rows).ok());

    const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
    EXPECT_EQ(result.get_null_map_data(), NullMap({0, 0, 0}));
    EXPECT_EQ(assert_cast<const ColumnInt32&>(result.get_nested_column()).get_data(),
              PaddedPODArray<Int32>({2, 3, 0}));
}

TEST(FunctionJsonbTEST, JsonLengthConstJsonVaryingPathMultiRows) {
    constexpr size_t rows = 3;
    auto jsonb_type = std::make_shared<DataTypeJsonb>();
    auto path_type = std::make_shared<DataTypeString>();
    auto return_type = make_nullable(std::make_shared<DataTypeInt32>());

    auto jsonb_column = jsonb_type->create_column();
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[1,2]})")));

    Block block;
    block.insert({ColumnConst::create(std::move(jsonb_column), rows), jsonb_type, "jsonb"});
    block.insert({create_path_column({"$.a", "$.a[0]", "$.missing"}), path_type, "path"});

    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            "json_length", block.get_columns_with_type_and_name(), return_type);
    ASSERT_NE(func, nullptr);

    FunctionUtils fn_utils(return_type, {jsonb_type, path_type}, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL).ok());
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::THREAD_LOCAL).ok());

    block.insert({nullptr, return_type, "result"});
    ASSERT_TRUE(func->execute(fn_ctx, block, {0, 1}, 2, rows).ok());

    const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
    EXPECT_EQ(result.get_null_map_data(), NullMap({0, 0, 1}));
    EXPECT_EQ(assert_cast<const ColumnInt32&>(result.get_nested_column()).get_data(),
              PaddedPODArray<Int32>({2, 1, 0}));
}

TEST(FunctionJsonbTEST, JsonContains) {
    std::string func_name = "json_contains";

    {
        InputTypeSet input_types = {PrimitiveType::TYPE_JSONB, PrimitiveType::TYPE_JSONB,
                                    PrimitiveType::TYPE_VARCHAR};

        DataSet data_set = {
                {{STRING(R"({"k1":"v1", "k2": 2})"), Null(), Null()}, Null()},
                {{STRING(R"([1, 2, 3])"), STRING("1"), STRING("$")}, BOOLEAN(true)},
                {{STRING(R"("string")"), STRING("string"), STRING("$")}, Null()},
                {{STRING(R"("string")"), STRING(R"("string")"), STRING("$")}, BOOLEAN(true)},
        };
        static_cast<void>(check_function<DataTypeBool, true>(func_name, input_types, data_set));
    }

    auto json_data_type = std::make_shared<DataTypeJsonb>();
    auto json_column = json_data_type->create_column();

    JsonbWriter writer;
    writer.writeStartArray();
    writer.writeString("hello");
    writer.writeString("world");
    writer.writeEndArray();

    json_column->insert_data(writer.getOutput()->getBuffer(), writer.getOutput()->getSize());

    auto json_column2 = json_data_type->create_column();
    writer.reset();
    writer.writeStartString();
    writer.writeString("hello");
    writer.writeEndString();

    Block block;
    block.insert(ColumnWithTypeAndName(std::move(json_column), json_data_type, "json_col"));
    block.insert(ColumnWithTypeAndName(std::move(json_column2), json_data_type, "json_col2"));

    auto path_column_ = ColumnNullable::create(ColumnString::create(), ColumnUInt8::create());
    path_column_->insert_default();
    auto path_column = ColumnConst::create(std::move(path_column_), 1);
    block.insert(ColumnWithTypeAndName(
            std::move(path_column),
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeString>()), "path_col"));

    auto return_type = make_nullable(std::make_shared<DataTypeBool>());
    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            func_name, block.get_columns_with_type_and_name(), return_type);
    ASSERT_TRUE(func != nullptr);

    std::vector<DataTypePtr> arg_types;
    arg_types.emplace_back(json_data_type);
    arg_types.emplace_back(json_data_type);
    arg_types.emplace_back(std::make_shared<DataTypeNullable>(std::make_shared<DataTypeString>()));

    FunctionUtils fn_utils(return_type, arg_types, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    //     fn_ctx->set_constant_cols(constant_cols);
    auto st = func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL);
    ASSERT_TRUE(st.ok()) << "open failed: " << st.to_string();
    st = func->open(fn_ctx, FunctionContext::THREAD_LOCAL);
    ASSERT_TRUE(st.ok()) << "open failed: " << st.to_string();

    block.insert({nullptr, return_type, "result"});

    auto result = block.columns() - 1;
    st = func->execute(fn_ctx, block, {0, 1, 2}, result, 1);
    ASSERT_TRUE(st.ok()) << "execute failed: " << st.to_string();
}

TEST(FunctionJsonbTEST, JsonContainsConstNullablePathMultiRows) {
    constexpr size_t rows = 3;
    auto jsonb_type = std::make_shared<DataTypeJsonb>();
    auto path_type = make_nullable(std::make_shared<DataTypeString>());
    auto return_type = make_nullable(std::make_shared<DataTypeBool>());

    auto jsonb_column = jsonb_type->create_column();
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[1,2]})")));
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[3,4,5]})")));
    ASSERT_TRUE(insert_cell(jsonb_column, jsonb_type, STRING(R"({"a":[]})")));

    auto candidate_column = jsonb_type->create_column();
    ASSERT_TRUE(insert_cell(candidate_column, jsonb_type, STRING("[3]")));

    Block block;
    block.insert({std::move(jsonb_column), jsonb_type, "jsonb"});
    block.insert({ColumnConst::create(std::move(candidate_column), rows), jsonb_type, "candidate"});
    block.insert({create_const_nullable_path_column("$.a", rows), path_type, "path"});

    FunctionBasePtr func = SimpleFunctionFactory::instance().get_function(
            "json_contains", block.get_columns_with_type_and_name(), return_type);
    ASSERT_NE(func, nullptr);

    FunctionUtils fn_utils(return_type, {jsonb_type, jsonb_type, path_type}, false);
    auto* fn_ctx = fn_utils.get_fn_ctx();
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL).ok());
    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::THREAD_LOCAL).ok());

    block.insert({nullptr, return_type, "result"});
    ASSERT_TRUE(func->execute(fn_ctx, block, {0, 1, 2}, 3, rows).ok());

    const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(3).column);
    EXPECT_EQ(result.get_null_map_data(), NullMap({0, 0, 0}));
    EXPECT_EQ(assert_cast<const ColumnUInt8&>(result.get_nested_column()).get_data(),
              PaddedPODArray<UInt8>({0, 1, 0}));
}

TEST(FunctionJsonbTEST, JsonbModifyMissingPathParent) {
    InputTypeSet input_types = {PrimitiveType::TYPE_JSONB, PrimitiveType::TYPE_VARCHAR,
                                PrimitiveType::TYPE_JSONB};

    {
        DataSet data_set = {
                {{STRING("{}"), STRING("$.a"), STRING("1")}, STRING(R"({"a":1})")},
                {{STRING(R"({"a":{}})"), STRING("$.a.b"), STRING("2")}, STRING(R"({"a":{"b":2}})")},
                {{STRING(R"({"a":[1]})"), STRING("$.a[1]"), STRING("2")}, STRING(R"({"a":[1,2]})")},
        };
        static_cast<void>(
                check_function<DataTypeJsonb, true>("jsonb_insert", input_types, data_set));
    }

    {
        DataSet data_set = {
                {{STRING("{}"), STRING("$.a"), STRING("1")}, STRING(R"({"a":1})")},
                {{STRING(R"({"a":{}})"), STRING("$.a.b"), STRING("2")}, STRING(R"({"a":{"b":2}})")},
                {{STRING(R"({"a":[1]})"), STRING("$.a[1]"), STRING("2")}, STRING(R"({"a":[1,2]})")},
        };
        static_cast<void>(check_function<DataTypeJsonb, true>("jsonb_set", input_types, data_set));
    }
}

TEST(FunctionJsonbTEST, JsonbModifyMemberPathOnArray) {
    InputTypeSet input_types = {PrimitiveType::TYPE_JSONB, PrimitiveType::TYPE_VARCHAR,
                                PrimitiveType::TYPE_JSONB};

    // A member leg never addresses an array element, so a missing member under an array is
    // not inserted; array legs and quoted literal keys still work.
    DataSet data_set = {
            {{STRING("[1]"), STRING("$.a"), STRING("2")}, STRING("[1]")},
            {{STRING("[]"), STRING("$.a"), STRING("2")}, STRING("[]")},
            {{STRING("[1]"), STRING(R"($."a.b")"), STRING("2")}, STRING("[1]")},
            {{STRING("[[1]]"), STRING("$[0].a"), STRING("2")}, STRING("[[1]]")},
            {{STRING(R"({"a":[1]})"), STRING("$.a.b"), STRING("2")}, STRING(R"({"a":[1]})")},
            {{STRING("[1]"), STRING("$[1]"), STRING("2")}, STRING("[1,2]")},
            {{STRING("{}"), STRING(R"($."a.b")"), STRING("2")}, STRING(R"({"a.b":2})")},
            {{STRING("{}"), STRING(R"($."")"), STRING("2")}, STRING(R"({"":2})")},
    };
    static_cast<void>(check_function<DataTypeJsonb, true>("jsonb_insert", input_types, data_set));
    static_cast<void>(check_function<DataTypeJsonb, true>("jsonb_set", input_types, data_set));

    // Quoted "*" and "" name existing literal keys, including as parents of a nested leg.
    DataSet set_existing = {
            {{STRING(R"({"*":1,"b":1})"), STRING(R"($."*")"), STRING("2")},
             STRING(R"({"*":2,"b":1})")},
            {{STRING(R"({"":1,"b":1})"), STRING(R"($."")"), STRING("2")},
             STRING(R"({"":2,"b":1})")},
            {{STRING(R"({"*":{"":1},"b":1})"), STRING(R"($."*"."")"), STRING("2")},
             STRING(R"({"*":{"":2},"b":1})")},
    };
    static_cast<void>(check_function<DataTypeJsonb, true>("jsonb_set", input_types, set_existing));
    static_cast<void>(
            check_function<DataTypeJsonb, true>("jsonb_replace", input_types, set_existing));

    // jsonb_insert keeps existing literal keys unchanged and only fills a missing nested key.
    DataSet insert_existing = {
            {{STRING(R"({"*":1,"b":1})"), STRING(R"($."*")"), STRING("2")},
             STRING(R"({"*":1,"b":1})")},
            {{STRING(R"({"":1,"b":1})"), STRING(R"($."")"), STRING("2")},
             STRING(R"({"":1,"b":1})")},
            {{STRING(R"({"*":{"":1},"b":1})"), STRING(R"($."*"."")"), STRING("2")},
             STRING(R"({"*":{"":1},"b":1})")},
            {{STRING(R"({"b":1,"*":{}})"), STRING(R"($."*"."")"), STRING("2")},
             STRING(R"({"b":1,"*":{"":2}})")},
    };
    static_cast<void>(
            check_function<DataTypeJsonb, true>("jsonb_insert", input_types, insert_existing));

    // jsonb_replace leaves a missing literal key alone.
    DataSet replace_missing = {
            {{STRING(R"({"b":1})"), STRING(R"($."*")"), STRING("2")}, STRING(R"({"b":1})")},
            {{STRING(R"({"b":1})"), STRING(R"($."")"), STRING("2")}, STRING(R"({"b":1})")},
    };
    static_cast<void>(
            check_function<DataTypeJsonb, true>("jsonb_replace", input_types, replace_missing));

    // An unquoted '*' is still a wildcard and is rejected by every modify function.
    DataSet wildcard = {
            {{STRING(R"({"*":1,"b":1})"), STRING("$.*"), STRING("2")}, Null()},
    };
    for (const auto* func_name : {"jsonb_insert", "jsonb_set", "jsonb_replace"}) {
        static_cast<void>(check_function<DataTypeJsonb, true>(func_name, input_types, wildcard, -1,
                                                              -1, true));
    }
}

} // namespace doris