// 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 <glog/logging.h>

#include <algorithm>

#include <cstdlib>

#include <memory>

#include <string>

#include <string_view>

#include <tuple>

#include <type_traits>

#include <utility>

#include <variant>

#include "common/compiler_util.h" // IWYU pragma: keep

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/block/column_numbers.h"

#include "core/block/column_with_type_and_name.h"

#include "core/column/column.h"

#include "core/column/column_array.h"

#include "core/column/column_const.h"

#include "core/column/column_nullable.h"

#include "core/column/column_string.h"

#include "core/column/column_vector.h"

#include "core/custom_allocator.h"

#include "core/data_type/data_type.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_string.h"

#include "core/data_type/define_primitive_type.h"

#include "core/data_type/primitive_type.h"

#include "core/string_ref.h"

#include "core/types.h"

#include "core/value/jsonb_value.h"

#include "exec/common/stringop_substring.h"

#include "exec/common/template_helpers.hpp"

#include "exec/common/util.hpp"

#include "exprs/aggregate/aggregate_function.h"

#include "exprs/function/function.h"

#include "exprs/function/like.h"

#include "exprs/function/simple_function_factory.h"

#include "exprs/function_context.h"

#include "util/jsonb_document.h"

#include "util/jsonb_utils.h"

#include "util/jsonb_writer.h"

#include "util/simd/bits.h"

namespace doris {

enum class NullalbeMode { NULLABLE = 0, FOLLOW_INPUT };

enum class JsonbParseErrorMode { FAIL = 0, RETURN_NULL, RETURN_VALUE };

// func(string,string) -> json

template <NullalbeMode nullable_mode, JsonbParseErrorMode parse_error_handle_mode>

class FunctionJsonbParseBase : public IFunction {

private:

struct FunctionJsonbParseState {

StringRef default_value;

JsonBinaryValue default_value_parser;

bool has_const_default_value = false;

bool default_is_null = false;

};

public:

static constexpr auto name = "json_parse";

static constexpr auto alias = "jsonb_parse";

static FunctionPtr create() { return std::make_shared<FunctionJsonbParseBase>(); }

String get_name() const override {

String error_mode;

switch (parse_error_handle_mode) {

case JsonbParseErrorMode::FAIL:

break;

case JsonbParseErrorMode::RETURN_NULL:

error_mode = "_error_to_null";

break;

case JsonbParseErrorMode::RETURN_VALUE:

error_mode = "_error_to_value";

break;

}

return name + error_mode;

}

bool is_variadic() const override {

return parse_error_handle_mode == JsonbParseErrorMode::RETURN_VALUE;

}

size_t get_number_of_arguments() const override {

switch (parse_error_handle_mode) {

case JsonbParseErrorMode::FAIL:

return 1;

case JsonbParseErrorMode::RETURN_NULL:

return 1;

case JsonbParseErrorMode::RETURN_VALUE:

return 0;

}

}

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

bool is_nullable = false;

switch (nullable_mode) {

case NullalbeMode::NULLABLE:

is_nullable = true;

break;

case NullalbeMode::FOLLOW_INPUT: {

for (auto arg : arguments) {

is_nullable |= arg->is_nullable();

}

break;

}

}

return is_nullable ? make_nullable(std::make_shared<DataTypeJsonb>())

: std::make_shared<DataTypeJsonb>();

}

bool use_default_implementation_for_nulls() const override { return false; }

Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {

if (scope == FunctionContext::FunctionStateScope::FRAGMENT_LOCAL) {

std::shared_ptr<FunctionJsonbParseState> state =

std::make_shared<FunctionJsonbParseState>();

context->set_function_state(FunctionContext::FRAGMENT_LOCAL, state);

}

if constexpr (parse_error_handle_mode == JsonbParseErrorMode::RETURN_VALUE) {

if (scope == FunctionContext::FunctionStateScope::FRAGMENT_LOCAL) {

auto* state = reinterpret_cast<FunctionJsonbParseState*>(

context->get_function_state(FunctionContext::FRAGMENT_LOCAL));

if (state) {

if (context->get_num_args() == 2) {

if (context->is_col_constant(1)) {

const auto default_value_col = context->get_constant_col(1)->column_ptr;

if (default_value_col->is_null_at(0)) {

state->default_is_null = true;

} else {

const auto& default_value = default_value_col->get_data_at(0);

state->default_value = default_value;

state->has_const_default_value = true;

}

}

} else if (context->get_num_args() == 1) {

RETURN_IF_ERROR(

state->default_value_parser.from_json_string(std::string("{}")));

state->default_value = StringRef(state->default_value_parser.value(),

state->default_value_parser.size());

state->has_const_default_value = true;

}

}

}

if (context->get_num_args() != 1 && context->get_num_args() != 2) {

return Status::InvalidArgument(

"{} function should have 1 or 2 arguments, "

"but got {}",

get_name(), context->get_num_args());

}

}

return Status::OK();

}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

auto&& [col_from, col_from_is_const] =

unpack_if_const(block.get_by_position(arguments[0]).column);

if (col_from_is_const && col_from->is_null_at(0)) {

auto col_str = ColumnString::create();

col_str->insert_default();

auto null_map = ColumnUInt8::create(1, 1);

auto nullable_col = ColumnNullable::create(std::move(col_str), std::move(null_map));

block.get_by_position(result).column =

ColumnConst::create(std::move(nullable_col), input_rows_count);

return Status::OK();

}

auto null_map = ColumnUInt8::create(0, 0);

bool is_nullable = false;

switch (nullable_mode) {

case NullalbeMode::NULLABLE: {

is_nullable = true;

break;

}

case NullalbeMode::FOLLOW_INPUT: {

for (auto arg : arguments) {

is_nullable |= block.get_by_position(arg).type->is_nullable();

}

break;

}

}

if (is_nullable) {

null_map = ColumnUInt8::create(input_rows_count, 0);

}

const ColumnString* col_from_string = nullptr;

if (const auto* nullable_col = check_and_get_column<ColumnNullable>(col_from.get())) {

VectorizedUtils::update_null_map(null_map->get_data(),

nullable_col->get_null_map_data(), col_from_is_const);

col_from_string =

assert_cast<const ColumnString*>(nullable_col->get_nested_column_ptr().get());

} else {

col_from_string = assert_cast<const ColumnString*>(col_from.get());

}

StringRef constant_default_value;

bool default_value_const = false;

bool default_value_null_const = false;

ColumnPtr default_value_col;

JsonBinaryValue default_jsonb_value_parser;

const ColumnString* default_value_str_col = nullptr;

const NullMap* default_value_nullmap = nullptr;

if constexpr (parse_error_handle_mode == JsonbParseErrorMode::RETURN_VALUE) {

auto* state = reinterpret_cast<FunctionJsonbParseState*>(

context->get_function_state(FunctionContext::FRAGMENT_LOCAL));

if (state && state->has_const_default_value) {

constant_default_value = state->default_value;

default_value_null_const = state->default_is_null;

default_value_const = true;

} else if (arguments.size() > 1) {

if (block.get_by_position(arguments[1]).type->get_primitive_type() !=

PrimitiveType::TYPE_JSONB) {

return Status::InvalidArgument(

"{} second argument should be jsonb type, but got {}", get_name(),

block.get_by_position(arguments[1]).type->get_name());

}

std::tie(default_value_col, default_value_const) =

unpack_if_const(block.get_by_position(arguments[1]).column);

if (default_value_const) {

const JsonbDocument* default_value_doc = nullptr;

if (default_value_col->is_null_at(0)) {

default_value_null_const = true;

} else {

auto data = default_value_col->get_data_at(0);

RETURN_IF_ERROR(JsonbDocument::checkAndCreateDocument(data.data, data.size,

&default_value_doc));

constant_default_value = data;

}

} else {

if (const auto* nullable_col =

check_and_get_column<ColumnNullable>(default_value_col.get())) {

default_value_str_col = assert_cast<const ColumnString*>(

nullable_col->get_nested_column_ptr().get());

default_value_nullmap = &(nullable_col->get_null_map_data());

} else {

default_value_str_col =

assert_cast<const ColumnString*>(default_value_col.get());

}

}

} else if (arguments.size() == 1) {

// parse default value '{}' should always success.

RETURN_IF_ERROR(default_jsonb_value_parser.from_json_string(std::string("{}")));

default_value_const = true;

constant_default_value.data = default_jsonb_value_parser.value();

constant_default_value.size = default_jsonb_value_parser.size();

}

}

auto col_to = ColumnString::create();

col_to->reserve(input_rows_count);

auto& null_map_data = null_map->get_data();

// parser can be reused for performance

JsonBinaryValue jsonb_value;

for (size_t i = 0; i < input_rows_count; ++i) {

if (is_nullable && null_map_data[i]) {

col_to->insert_default();

continue;

}

auto index = index_check_const(i, col_from_is_const);

const auto& val = col_from_string->get_data_at(index);

auto st = jsonb_value.from_json_string(val.data, val.size);

if (st.ok()) {

// insert jsonb format data

col_to->insert_data(jsonb_value.value(), jsonb_value.size());

} else {

if constexpr (parse_error_handle_mode == JsonbParseErrorMode::FAIL) {

return Status::InvalidArgument(

"Parse json document failed at row {}, error: {}", i, st.to_string());

} else if constexpr (parse_error_handle_mode == JsonbParseErrorMode::RETURN_NULL) {

null_map_data[i] = 1;

col_to->insert_default();

} else {

if (default_value_const) {

if (default_value_null_const) {

null_map_data[i] = 1;

col_to->insert_default();

} else {

col_to->insert_data(constant_default_value.data,

constant_default_value.size);

}

} else {

if (default_value_nullmap && (*default_value_nullmap)[i]) {

null_map_data[i] = 1;

col_to->insert_default();

continue;

}

auto value = default_value_str_col->get_data_at(i);

col_to->insert_data(value.data, value.size);

}

}

}

}

if (is_nullable) {

block.replace_by_position(

result, ColumnNullable::create(std::move(col_to), std::move(null_map)));

} else {

block.replace_by_position(result, std::move(col_to));

}

return Status::OK();

}

};

// jsonb_parse return type nullable as input

using FunctionJsonbParse =

FunctionJsonbParseBase<NullalbeMode::FOLLOW_INPUT, JsonbParseErrorMode::FAIL>;

using FunctionJsonbParseErrorNull =

FunctionJsonbParseBase<NullalbeMode::NULLABLE, JsonbParseErrorMode::RETURN_NULL>;

using FunctionJsonbParseErrorValue =

FunctionJsonbParseBase<NullalbeMode::FOLLOW_INPUT, JsonbParseErrorMode::RETURN_VALUE>;

// func(jsonb, [varchar, varchar, ...]) -> nullable(type)

template <typename Impl>

class FunctionJsonbExtract : public IFunction {

public:

static constexpr auto name = Impl::name;

static constexpr auto alias = Impl::alias;

static FunctionPtr create() { return std::make_shared<FunctionJsonbExtract>(); }

String get_name() const override { return name; }

bool is_variadic() const override { return true; }

size_t get_number_of_arguments() const override { return 0; }

bool use_default_implementation_for_nulls() const override { return false; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

return make_nullable(std::make_shared<typename Impl::ReturnType>());

}

DataTypes get_variadic_argument_types_impl() const override {

if constexpr (HasGetVariadicArgumentTypesImpl<Impl>) {

return Impl::get_variadic_argument_types_impl();

} else {

return {};

}

}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

DORIS_CHECK_GE(arguments.size(), 2);

ColumnPtr jsonb_data_column;

bool jsonb_data_const = false;

const NullMap* data_null_map = nullptr;

if (block.get_by_position(arguments[0]).type->get_primitive_type() !=

PrimitiveType::TYPE_JSONB) {

return Status::InvalidArgument(

"jsonb_extract first argument should be json type, but got {}",

block.get_by_position(arguments[0]).type->get_name());

}

// prepare jsonb data column

std::tie(jsonb_data_column, jsonb_data_const) =

unpack_if_const(block.get_by_position(arguments[0]).column);

if (const auto* nullable_column =

check_and_get_column<ColumnNullable>(jsonb_data_column.get())) {

jsonb_data_column = nullable_column->get_nested_column_ptr();

data_null_map = &nullable_column->get_null_map_data();

}

const auto& ldata = assert_cast<const ColumnString*>(jsonb_data_column.get())->get_chars();

const auto& loffsets =

assert_cast<const ColumnString*>(jsonb_data_column.get())->get_offsets();

// prepare parse path column prepare

std::vector<const ColumnString*> jsonb_path_columns;

std::vector<bool> path_const(arguments.size() - 1);

std::vector<const NullMap*> path_null_maps(arguments.size() - 1, nullptr);

for (int i = 0; i < arguments.size() - 1; ++i) {

ColumnPtr path_column;

bool is_const = false;

std::tie(path_column, is_const) =

unpack_if_const(block.get_by_position(arguments[i + 1]).column);

path_const[i] = is_const;

if (const auto* nullable_column =

check_and_get_column<ColumnNullable>(path_column.get())) {

path_column = nullable_column->get_nested_column_ptr();

path_null_maps[i] = &nullable_column->get_null_map_data();

}

jsonb_path_columns.push_back(assert_cast<const ColumnString*>(path_column.get()));

}

auto null_map = ColumnUInt8::create(input_rows_count, 0);

auto res = Impl::ColumnType::create();

// execute Impl

if constexpr (std::is_same_v<typename Impl::ReturnType, DataTypeString> ||

std::is_same_v<typename Impl::ReturnType, DataTypeJsonb>) {

auto& res_data = res->get_chars();

auto& res_offsets = res->get_offsets();

RETURN_IF_ERROR(Impl::vector_vector_v2(

context, ldata, loffsets, data_null_map, jsonb_data_const, jsonb_path_columns,

path_null_maps, path_const, res_data, res_offsets, null_map->get_data()));

} else {

// not support other extract type for now (e.g. int, double, ...)

DORIS_CHECK_EQ(jsonb_path_columns.size(), 1);

const auto& rdata = jsonb_path_columns[0]->get_chars();

const auto& roffsets = jsonb_path_columns[0]->get_offsets();

auto create_all_null_result = [&]() {

res = Impl::ColumnType::create();

res->insert_default();

auto nullable_column =

ColumnNullable::create(std::move(res), ColumnUInt8::create(1, 1));

auto const_column =

ColumnConst::create(std::move(nullable_column), input_rows_count);

block.get_by_position(result).column = std::move(const_column);

return Status::OK();

};

if (jsonb_data_const) {

if (data_null_map && (*data_null_map)[0]) {

return create_all_null_result();

}

RETURN_IF_ERROR(Impl::scalar_vector(context, jsonb_data_column->get_data_at(0),

rdata, roffsets, path_null_maps[0],

res->get_data(), null_map->get_data()));

} else if (path_const[0]) {

if (path_null_maps[0] && (*path_null_maps[0])[0]) {

return create_all_null_result();

}

RETURN_IF_ERROR(Impl::vector_scalar(context, ldata, loffsets, data_null_map,

jsonb_path_columns[0]->get_data_at(0),

res->get_data(), null_map->get_data()));

} else {

RETURN_IF_ERROR(Impl::vector_vector(context, ldata, loffsets, data_null_map, rdata,

roffsets, path_null_maps[0], res->get_data(),

null_map->get_data()));

}

}

block.get_by_position(result).column =

ColumnNullable::create(std::move(res), std::move(null_map));

return Status::OK();

}

};

class FunctionJsonbKeys : public IFunction {

public:

static constexpr auto name = "json_keys";

static constexpr auto alias = "jsonb_keys";

static FunctionPtr create() { return std::make_shared<FunctionJsonbKeys>(); }

String get_name() const override { return name; }

bool is_variadic() const override { return true; }

size_t get_number_of_arguments() const override { return 0; }

bool use_default_implementation_for_nulls() const override { return false; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

return make_nullable(

std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));

}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

DORIS_CHECK_GE(arguments.size(), 1);

DORIS_CHECK(arguments.size() == 1 || arguments.size() == 2)

<< "json_keys should have 1 or 2 arguments, but got " << arguments.size();

const NullMap* data_null_map = nullptr;

const ColumnString* col_from_string = nullptr;

// prepare jsonb data column

auto&& [jsonb_data_column, json_data_const] =

unpack_if_const(block.get_by_position(arguments[0]).column);

if (const auto* nullable = check_and_get_column<ColumnNullable>(jsonb_data_column.get())) {

col_from_string =

assert_cast<const ColumnString*>(nullable->get_nested_column_ptr().get());

data_null_map = &nullable->get_null_map_data();

} else {

col_from_string = assert_cast<const ColumnString*>(jsonb_data_column.get());

}

// prepare parse path column prepare, maybe we do not have path column

ColumnPtr jsonb_path_column = nullptr;

const ColumnString* jsonb_path_col = nullptr;

bool path_const = false;

const NullMap* path_null_map = nullptr;

if (arguments.size() == 2) {

// we have should have a ColumnString for path

std::tie(jsonb_path_column, path_const) =

unpack_if_const(block.get_by_position(arguments[1]).column);

if (const auto* nullable =

check_and_get_column<ColumnNullable>(jsonb_path_column.get())) {

jsonb_path_column = nullable->get_nested_column_ptr();

path_null_map = &nullable->get_null_map_data();

}

jsonb_path_col = check_and_get_column<ColumnString>(jsonb_path_column.get());

}

auto null_map = ColumnUInt8::create(input_rows_count, 0);

NullMap& res_null_map = null_map->get_data();

auto dst_arr = ColumnArray::create(

ColumnNullable::create(ColumnString::create(), ColumnUInt8::create()),

ColumnArray::ColumnOffsets::create());

auto& dst_nested_column = assert_cast<ColumnNullable&>(dst_arr->get_data());

Status st = std::visit(

[&](auto data_const, auto has_path, auto path_const) {

return inner_loop_impl<data_const, has_path, path_const>(

input_rows_count, *dst_arr, dst_nested_column, res_null_map,

*col_from_string, data_null_map, jsonb_path_col, path_null_map);

},

make_bool_variant(json_data_const), make_bool_variant(jsonb_path_column),

make_bool_variant(path_const));

if (!st.ok()) {

return st;

}

block.get_by_position(result).column =

ColumnNullable::create(std::move(dst_arr), std::move(null_map));

return st;

}

private:

template <bool JSONB_DATA_CONST, bool JSONB_PATH_PARAM, bool JSON_PATH_CONST>

static ALWAYS_INLINE Status inner_loop_impl(size_t input_rows_count, ColumnArray& dst_arr,

ColumnNullable& dst_nested_column,

NullMap& res_null_map,

const ColumnString& col_from_string,

const NullMap* jsonb_data_nullmap,

const ColumnString* jsonb_path_column,

const NullMap* path_null_map) {

// if path is const, we just need to parse it once

JsonbPath const_path;

if constexpr (JSONB_PATH_PARAM && JSON_PATH_CONST) {

StringRef r_raw_ref = jsonb_path_column->get_data_at(0);

if (!const_path.seek(r_raw_ref.data, r_raw_ref.size)) {

return Status::InvalidArgument("Json path error: Invalid Json Path for value: {}",

r_raw_ref.to_string());

}

if (const_path.is_wildcard() || const_path.is_supper_wildcard()) {

return Status::InvalidJsonPath(

"In this situation, path expressions may not contain the * and ** tokens "

"or an array range.");

}

}

for (size_t i = 0; i < input_rows_count; ++i) {

auto index = index_check_const(i, JSONB_DATA_CONST);

// if jsonb data is null or path column is null , we should return null

if (jsonb_data_nullmap && (*jsonb_data_nullmap)[index]) {

res_null_map[i] = 1;

dst_arr.insert_default();

continue;

}

if constexpr (JSONB_PATH_PARAM && !JSON_PATH_CONST) {

if (path_null_map && (*path_null_map)[i]) {

res_null_map[i] = 1;

dst_arr.insert_default();

continue;

}

}

auto json_data = col_from_string.get_data_at(index);

const JsonbDocument* doc = nullptr;

auto st = JsonbDocument::checkAndCreateDocument(json_data.data, json_data.size, &doc);

if (!st.ok() || !doc || !doc->getValue()) [[unlikely]] {

dst_arr.clear();

return Status::InvalidArgument("jsonb data is invalid");

}

const JsonbValue* obj_val;

JsonbFindResult find_result;

if constexpr (JSONB_PATH_PARAM) {

if constexpr (!JSON_PATH_CONST) {

auto data = jsonb_path_column->get_data_at(i);

JsonbPath path;

if (!path.seek(data.data, data.size)) {

return Status::InvalidArgument(

"Json path error: Invalid Json Path for value: {} at row: {}",

std::string_view(data.data, data.size), i);

}

if (path.is_wildcard() || path.is_supper_wildcard()) {

return Status::InvalidJsonPath(

"In this situation, path expressions may not contain the * and ** "

"tokens "

"or an array range. at row: {}",

i);

}

find_result = doc->getValue()->findValue(path);

} else {

find_result = doc->getValue()->findValue(const_path);

}

obj_val = find_result.value;

} else {

obj_val = doc->getValue();

}

if (!obj_val || !obj_val->isObject()) {

// if jsonb data is not object we should return null

res_null_map[i] = 1;

dst_arr.insert_default();

continue;

}

const auto* obj = obj_val->unpack<ObjectVal>();

for (const auto& it : *obj) {

dst_nested_column.insert_data(it.getKeyStr(), it.klen());

}

dst_arr.get_offsets().push_back(dst_nested_column.size());

} //for

return Status::OK();

}

};

class FunctionJsonbExtractPath : public IFunction {

public:

static constexpr auto name = "json_exists_path";

static constexpr auto alias = "jsonb_exists_path";

using ColumnType = ColumnUInt8;

using Container = typename ColumnType::Container;

static FunctionPtr create() { return std::make_shared<FunctionJsonbExtractPath>(); }

String get_name() const override { return name; }

size_t get_number_of_arguments() const override { return 2; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

// it only needs to indicate existence and does not need to return nullable values.

const auto nullable = std::ranges::any_of(

arguments, [](const DataTypePtr& type) { return type->is_nullable(); });

if (nullable) {

return make_nullable(std::make_shared<DataTypeUInt8>());

} else {

return std::make_shared<DataTypeUInt8>();

}

}

bool use_default_implementation_for_nulls() const override { return false; }

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

// prepare jsonb data column

auto&& [jsonb_data_column, jsonb_data_const] =

unpack_if_const(block.get_by_position(arguments[0]).column);

const NullMap* data_null_map = nullptr;

const ColumnString* data_col = nullptr;

if (const auto* nullable = check_and_get_column<ColumnNullable>(jsonb_data_column.get())) {

data_col = assert_cast<const ColumnString*>(nullable->get_nested_column_ptr().get());

data_null_map = &nullable->get_null_map_data();

} else {

data_col = assert_cast<const ColumnString*>(jsonb_data_column.get());

}

const auto& ldata = data_col->get_chars();

const auto& loffsets = data_col->get_offsets();

// prepare parse path column prepare

auto&& [path_column, path_const] =

unpack_if_const(block.get_by_position(arguments[1]).column);

const ColumnString* path_col = nullptr;

const NullMap* path_null_map = nullptr;

if (const auto* nullable = check_and_get_column<ColumnNullable>(path_column.get())) {

path_col = assert_cast<const ColumnString*>(nullable->get_nested_column_ptr().get());

path_null_map = &nullable->get_null_map_data();

} else {

path_col = assert_cast<const ColumnString*>(path_column.get());

}

DORIS_CHECK(!(jsonb_data_const && path_const))

<< "jsonb_data_const and path_const should not be both const";

auto create_all_null_result = [&]() {

auto res = ColumnType::create();

res->insert_default();

auto nullable_column =

ColumnNullable::create(std::move(res), ColumnUInt8::create(1, 1));

auto const_column = ColumnConst::create(std::move(nullable_column), input_rows_count);

block.get_by_position(result).column = std::move(const_column);

return Status::OK();

};

ColumnUInt8::MutablePtr result_null_map_column;

NullMap* result_null_map = nullptr;

if (data_null_map || path_null_map) {

result_null_map_column = ColumnUInt8::create(input_rows_count, 0);

result_null_map = &result_null_map_column->get_data();

if (data_null_map) {

VectorizedUtils::update_null_map(*result_null_map, *data_null_map,

jsonb_data_const);

}

if (path_null_map) {

VectorizedUtils::update_null_map(*result_null_map, *path_null_map, path_const);

}

if (!simd::contain_zero(result_null_map->data(), input_rows_count)) {

return create_all_null_result();

}

}

auto res = ColumnType::create();

bool is_invalid_json_path = false;

const auto& rdata = path_col->get_chars();

const auto& roffsets = path_col->get_offsets();

if (jsonb_data_const) {

if (data_null_map && (*data_null_map)[0]) {

return create_all_null_result();

}

scalar_vector(context, data_col->get_data_at(0), rdata, roffsets, res->get_data(),

result_null_map, is_invalid_json_path);

} else if (path_const) {

if (path_null_map && (*path_null_map)[0]) {

return create_all_null_result();

}

vector_scalar(context, ldata, loffsets, path_col->get_data_at(0), res->get_data(),

result_null_map, is_invalid_json_path);

} else {

vector_vector(context, ldata, loffsets, rdata, roffsets, res->get_data(),

result_null_map, is_invalid_json_path);

}

if (is_invalid_json_path) {

return Status::InvalidArgument(

"Json path error: Invalid Json Path for value: {}",

std::string_view(reinterpret_cast<const char*>(rdata.data()), rdata.size()));

}

if (result_null_map) {

auto nullabel_col =

ColumnNullable::create(std::move(res), std::move(result_null_map_column));

block.get_by_position(result).column = std::move(nullabel_col);

} else {

block.get_by_position(result).column = std::move(res);

}

return Status::OK();

}

private:

static ALWAYS_INLINE void inner_loop_impl(size_t i, Container& res, const char* l_raw_str,

size_t l_str_size, JsonbPath& path) {

// doc is NOT necessary to be deleted since JsonbDocument will not allocate memory

const JsonbDocument* doc = nullptr;

auto st = JsonbDocument::checkAndCreateDocument(l_raw_str, l_str_size, &doc);

if (!st.ok() || !doc || !doc->getValue()) [[unlikely]] {

return;

}

// value is NOT necessary to be deleted since JsonbValue will not allocate memory

auto result = doc->getValue()->findValue(path);

if (result.value) {

res[i] = 1;

}

}

static void vector_vector(FunctionContext* context, const ColumnString::Chars& ldata,

const ColumnString::Offsets& loffsets,

const ColumnString::Chars& rdata,

const ColumnString::Offsets& roffsets, Container& res,

const NullMap* result_null_map, bool& is_invalid_json_path) {

const size_t size = loffsets.size();

res.resize_fill(size, 0);

for (size_t i = 0; i < size; i++) {

if (result_null_map && (*result_null_map)[i]) {

continue;

}

const char* l_raw_str = reinterpret_cast<const char*>(&ldata[loffsets[i - 1]]);

int l_str_size = loffsets[i] - loffsets[i - 1];

const char* r_raw_str = reinterpret_cast<const char*>(&rdata[roffsets[i - 1]]);

int r_str_size = roffsets[i] - roffsets[i - 1];

JsonbPath path;

if (!path.seek(r_raw_str, r_str_size)) {

is_invalid_json_path = true;

return;

}

inner_loop_impl(i, res, l_raw_str, l_str_size, path);

}

}

static void scalar_vector(FunctionContext* context, const StringRef& ldata,

const ColumnString::Chars& rdata,

const ColumnString::Offsets& roffsets, Container& res,

const NullMap* result_null_map, bool& is_invalid_json_path) {

const size_t size = roffsets.size();

res.resize_fill(size, 0);

for (size_t i = 0; i < size; i++) {

if (result_null_map && (*result_null_map)[i]) {

continue;

}

const char* r_raw_str = reinterpret_cast<const char*>(&rdata[roffsets[i - 1]]);

int r_str_size = roffsets[i] - roffsets[i - 1];

JsonbPath path;

if (!path.seek(r_raw_str, r_str_size)) {

is_invalid_json_path = true;

return;

}

inner_loop_impl(i, res, ldata.data, ldata.size, path);

}

}

static void vector_scalar(FunctionContext* context, const ColumnString::Chars& ldata,

const ColumnString::Offsets& loffsets, const StringRef& rdata,

Container& res, const NullMap* result_null_map,

bool& is_invalid_json_path) {

const size_t size = loffsets.size();

res.resize_fill(size, 0);

JsonbPath path;

if (!path.seek(rdata.data, rdata.size)) {

is_invalid_json_path = true;

return;

}

for (size_t i = 0; i < size; i++) {

if (result_null_map && (*result_null_map)[i]) {

continue;

}

const char* l_raw_str = reinterpret_cast<const char*>(&ldata[loffsets[i - 1]]);

int l_str_size = loffsets[i] - loffsets[i - 1];

inner_loop_impl(i, res, l_raw_str, l_str_size, path);

}

}

};

template <typename ValueType>

struct JsonbExtractStringImpl {

using ReturnType = typename ValueType::ReturnType;

using ColumnType = typename ValueType::ColumnType;

private:

static ALWAYS_INLINE void inner_loop_impl(JsonbWriter* writer, size_t i,

ColumnString::Chars& res_data,

ColumnString::Offsets& res_offsets, NullMap& null_map,

const char* l_raw, size_t l_size, JsonbPath& path) {

// doc is NOT necessary to be deleted since JsonbDocument will not allocate memory

const JsonbDocument* doc = nullptr;

auto st = JsonbDocument::checkAndCreateDocument(l_raw, l_size, &doc);

if (!st.ok() || !doc || !doc->getValue()) [[unlikely]] {

StringOP::push_null_string(i, res_data, res_offsets, null_map);

return;

}

// value is NOT necessary to be deleted since JsonbValue will not allocate memory

auto find_result = doc->getValue()->findValue(path);

if (UNLIKELY(!find_result.value)) {

StringOP::push_null_string(i, res_data, res_offsets, null_map);

return;

}

if constexpr (ValueType::only_get_type) {

StringOP::push_value_string(std::string_view(find_result.value->typeName()), i,

res_data, res_offsets);

return;

} else {

static_assert(std::is_same_v<DataTypeJsonb, ReturnType>);

if constexpr (ValueType::no_quotes) {

if (find_result.value->isString()) {

const auto* str_value = find_result.value->unpack<JsonbStringVal>();

const auto* blob = str_value->getBlob();

if (str_value->length() > 1 && blob[0] == '"' &&

blob[str_value->length() - 1] == '"') {

writer->writeStartString();

writer->writeString(blob + 1, str_value->length() - 2);

writer->writeEndString();

StringOP::push_value_string(

std::string_view(writer->getOutput()->getBuffer(),

writer->getOutput()->getSize()),

i, res_data, res_offsets);

return;

}

}

}

writer->writeValueSimple(find_result.value);

StringOP::push_value_string(std::string_view(writer->getOutput()->getBuffer(),

writer->getOutput()->getSize()),

i, res_data, res_offsets);

}

}

public:

// for jsonb_extract_string

static Status vector_vector_v2(

FunctionContext* context, const ColumnString::Chars& ldata,

const ColumnString::Offsets& loffsets, const NullMap* l_null_map,

const bool& json_data_const,

const std::vector<const ColumnString*>& rdata_columns, // here we can support more paths

const std::vector<const NullMap*>& r_null_maps, const std::vector<bool>& path_const,

ColumnString::Chars& res_data, ColumnString::Offsets& res_offsets, NullMap& null_map) {

const size_t input_rows_count = null_map.size();

res_offsets.resize(input_rows_count);

auto writer = std::make_unique<JsonbWriter>();

// reuseable json path list, espacially for const path

std::vector<JsonbPath> json_path_list;

json_path_list.resize(rdata_columns.size());

// lambda function to parse json path for row i and path pi

auto parse_json_path = [&](size_t i, size_t pi) -> Status {

const auto index = index_check_const(i, path_const[pi]);

const ColumnString* path_col = rdata_columns[pi];

const ColumnString::Chars& rdata = path_col->get_chars();

const ColumnString::Offsets& roffsets = path_col->get_offsets();

size_t r_off = roffsets[index - 1];

size_t r_size = roffsets[index] - r_off;

const char* r_raw = reinterpret_cast<const char*>(&rdata[r_off]);

JsonbPath path;

if (!path.seek(r_raw, r_size)) {

return Status::InvalidArgument("Json path error: Invalid Json Path for value: {}",

std::string_view(r_raw, r_size));

}

json_path_list[pi] = std::move(path);

return Status::OK();

};

for (size_t pi = 0; pi < rdata_columns.size(); pi++) {

if (path_const[pi]) {

if (r_null_maps[pi] && (*r_null_maps[pi])[0]) {

continue;

}

RETURN_IF_ERROR(parse_json_path(0, pi));

}

}

res_data.reserve(ldata.size());

for (size_t i = 0; i < input_rows_count; ++i) {

if (null_map[i]) {

continue;

}

const auto data_index = index_check_const(i, json_data_const);

if (l_null_map && (*l_null_map)[data_index]) {

StringOP::push_null_string(i, res_data, res_offsets, null_map);

continue;

}

size_t l_off = loffsets[data_index - 1];

size_t l_size = loffsets[data_index] - l_off;

const char* l_raw = reinterpret_cast<const char*>(&ldata[l_off]);

if (rdata_columns.size() == 1) { // just return origin value

const auto path_index = index_check_const(i, path_const[0]);

if (r_null_maps[0] && (*r_null_maps[0])[path_index]) {

StringOP::push_null_string(i, res_data, res_offsets, null_map);

continue;

}

if (!path_const[0]) {

RETURN_IF_ERROR(parse_json_path(i, 0));

}

writer->reset();

inner_loop_impl(writer.get(), i, res_data, res_offsets, null_map, l_raw, l_size,

json_path_list[0]);

} else { // will make array string to user

writer->reset();

bool has_value = false;

// doc is NOT necessary to be deleted since JsonbDocument will not allocate memory

const JsonbDocument* doc = nullptr;

auto st = JsonbDocument::checkAndCreateDocument(l_raw, l_size, &doc);

for (size_t pi = 0; pi < rdata_columns.size(); ++pi) {

if (!st.ok() || !doc || !doc->getValue()) [[unlikely]] {

continue;