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

#include <vector>

#include "common/status.h"

#include "core/data_type/data_type_jsonb.h"

#include "core/data_type/primitive_type.h"

#include "exprs/function/simple_function_factory.h"

#include "util/jsonb_document.h"

#include "util/jsonb_document_cast.h"

#include "util/jsonb_writer.h"

namespace doris {

// Sort the keys of the JSON object and deduplicate the repeated keys, keeping the first one

void sort_json_object_keys(JsonbWriter& jsonb_writer, const JsonbValue* jsonb_value) {

if (jsonb_value->isObject()) {

std::vector<std::pair<StringRef, const JsonbValue*>> kvs;

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

for (auto it = obj_val->begin(); it != obj_val->end(); ++it) {

kvs.emplace_back(StringRef(it->getKeyStr(), it->klen()), it->value());

}

// sort by key

std::sort(kvs.begin(), kvs.end(),

[](const auto& left, const auto& right) { return left.first < right.first; });

// unique by key

kvs.erase(std::unique(kvs.begin(), kvs.end(),

[](const auto& left, const auto& right) {

return left.first == right.first;

}),

kvs.end());

jsonb_writer.writeStartObject();

for (const auto& kv : kvs) {

jsonb_writer.writeKey(kv.first.data, static_cast<uint8_t>(kv.first.size));

sort_json_object_keys(jsonb_writer, kv.second);

}

jsonb_writer.writeEndObject();

} else if (jsonb_value->isArray()) {

const auto* array_val = jsonb_value->unpack<ArrayVal>();

jsonb_writer.writeStartArray();

for (auto it = array_val->begin(); it != array_val->end(); ++it) {

sort_json_object_keys(jsonb_writer, &*it);

}

jsonb_writer.writeEndArray();

} else {

// scalar value

jsonb_writer.writeValue(jsonb_value);

}

}

// Walk a JSONB object recursively and emit flat "<dot.path>": value entries

// directly into `writer`. Members whose value is a non-empty object recurse;

// every other shape (scalars, arrays, null literals, empty objects) is emitted

// as an opaque leaf at its dot-joined path. The `prefix` buffer is reused

// across the whole row — appended on descent and truncated on return — so no

// path segment is ever re-allocated outside this single growing string.

void flatten_json_object_into(JsonbWriter& jsonb_writer, const ObjectVal* obj,

std::string& prefix) {

for (auto it = obj->begin(); it != obj->end(); ++it) {

const auto* val = it->value();

const size_t saved = prefix.size();

if (!prefix.empty()) {

prefix.push_back('.');

}

prefix.append(it->getKeyStr(), it->klen());

if (val->isObject() && val->unpack<ObjectVal>()->numElem() > 0) {

flatten_json_object_into(jsonb_writer, val->unpack<ObjectVal>(), prefix);

} else {

jsonb_writer.writeKey(prefix.data(), static_cast<uint8_t>(prefix.size()));

jsonb_writer.writeValue(val);

}

prefix.resize(saved);

}

}

// json_object_flatten: turn a nested JSONB object into a single-level JSONB

// object whose keys are the dot-joined paths to each leaf (NiFi FlattenJson

// "keep-arrays" semantics — arrays / scalars / nulls / empty objects stay as

// opaque leaf values; only objects are walked).

// {"a":{"b":2}} -> {"a.b":2}

// {"a":[{"b":1}]} -> {"a":[{"b":1}]}

// Top-level non-object values pass through unchanged.

void flatten_json_object(JsonbWriter& jsonb_writer, const JsonbValue* jsonb_value) {

if (!jsonb_value->isObject()) {

jsonb_writer.writeValue(jsonb_value);

return;

}

jsonb_writer.writeStartObject();

std::string prefix;

flatten_json_object_into(jsonb_writer, jsonb_value->unpack<ObjectVal>(), prefix);

jsonb_writer.writeEndObject();

}

// Convert all numeric types in JSON to double type

void normalize_json_numbers_to_double(JsonbWriter& jsonb_writer, const JsonbValue* jsonb_value) {

if (jsonb_value->isObject()) {

jsonb_writer.writeStartObject();

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

for (auto it = obj_val->begin(); it != obj_val->end(); ++it) {

jsonb_writer.writeKey(it->getKeyStr(), it->klen());

normalize_json_numbers_to_double(jsonb_writer, it->value());

}

jsonb_writer.writeEndObject();

} else if (jsonb_value->isArray()) {

const auto* array_val = jsonb_value->unpack<ArrayVal>();

jsonb_writer.writeStartArray();

for (auto it = array_val->begin(); it != array_val->end(); ++it) {

normalize_json_numbers_to_double(jsonb_writer, &*it);

}

jsonb_writer.writeEndArray();

} else {

// scalar value

if (jsonb_value->isInt() || jsonb_value->isFloat() || jsonb_value->isDouble() ||

jsonb_value->isDecimal()) {

double to;

CastParameters params;

params.is_strict = false;

JsonbCast::cast_from_json_to_float(jsonb_value, to, params);

NormalizeFloat(to);

jsonb_writer.writeDouble(to);

} else {

jsonb_writer.writeValue(jsonb_value);

}

}

}

// Input jsonb, output jsonb

template <typename Impl>

class FunctionJsonbTransform : public IFunction {

public:

static constexpr auto name = Impl::name;

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

String get_name() const override { return name; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

return std::make_shared<DataTypeJsonb>();

}

size_t get_number_of_arguments() const override { return 1; }

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

uint32_t result, size_t size) const override {

auto input_column = block.get_by_position(arguments[0]).column;

auto to_column = ColumnString::create();

const auto& input_jsonb_column = assert_cast<const ColumnString&>(*input_column);

to_column->get_chars().reserve(input_jsonb_column.get_chars().size());

to_column->get_offsets().reserve(input_jsonb_column.get_offsets().size());

JsonbWriter writer;

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

StringRef val = input_jsonb_column.get_data_at(i);

const JsonbDocument* doc = nullptr;

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

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

// mayby be invalid jsonb, just insert default

// invalid jsonb value may be caused by the default null processing

// insert empty string

to_column->insert_default();

continue;

}

const JsonbValue* value = doc->getValue();

if (UNLIKELY(!value)) {

// mayby be invalid jsonb, just insert default

// invalid jsonb value may be caused by the default null processing

// insert empty string

to_column->insert_default();

continue;

}

writer.reset();

Impl::transform(writer, value);

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

}

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

return Status::OK();

}

};

struct SortJsonObjectKeys {

static constexpr auto name = "sort_json_object_keys";

static void transform(JsonbWriter& writer, const JsonbValue* value) {

sort_json_object_keys(writer, value);

}

};

struct NormalizeJsonNumbersToDouble {

static constexpr auto name = "normalize_json_numbers_to_double";

static void transform(JsonbWriter& writer, const JsonbValue* value) {

normalize_json_numbers_to_double(writer, value);

}

};

struct JsonObjectFlatten {

static constexpr auto name = "json_object_flatten";

static void transform(JsonbWriter& writer, const JsonbValue* value) {

flatten_json_object(writer, value);

}

};

using FunctionSortJsonObjectKeys = FunctionJsonbTransform<SortJsonObjectKeys>;

using FunctionNormalizeJsonNumbersToDouble = FunctionJsonbTransform<NormalizeJsonNumbersToDouble>;

using FunctionJsonObjectFlatten = FunctionJsonbTransform<JsonObjectFlatten>;

void register_function_json_transform(SimpleFunctionFactory& factory) {

factory.register_function<FunctionSortJsonObjectKeys>();

factory.register_function<FunctionNormalizeJsonNumbersToDouble>();

factory.register_function<FunctionJsonObjectFlatten>();

factory.register_alias(FunctionSortJsonObjectKeys::name, "sort_jsonb_object_keys");

factory.register_alias(FunctionNormalizeJsonNumbersToDouble::name,

"normalize_jsonb_numbers_to_double");

}

} // namespace doris