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

#include <memory>

#include <string>

#include <vector>

#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/block/columns_with_type_and_name.h"

#include "core/column/column.h"

#include "core/column/column_array.h"

#include "core/column/column_nullable.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type.h"

#include "core/data_type/data_type_array.h"

#include "core/data_type/data_type_nullable.h"

#include "core/data_type/data_type_number.h"

#include "exec/common/util.hpp"

#include "exprs/aggregate/aggregate_function.h"

#include "exprs/lambda_function/lambda_function.h"

#include "exprs/lambda_function/lambda_function_factory.h"

#include "exprs/vcolumn_ref.h"

#include "exprs/vslot_ref.h"

namespace doris {

class VExprContext;

// extend a block with all required parameters

struct LambdaArgs {

// the lambda function need the column ids of all the slots

std::vector<int> output_slot_ref_indexs;

// which line is extended to the original block

int64_t current_row_idx = 0;

// when a block is filled, the array may be truncated, recording where it was truncated

int64_t current_offset_in_array = 0;

// the beginning position of the array

size_t array_start = 0;

// the size of the array

int64_t cur_size = 0;

// offset of column array

const ColumnArray::Offsets64* offsets_ptr = nullptr;

// expend data of repeat times

int current_repeat_times = 0;

// whether the current row of the original block has been extended

bool current_row_eos = false;

};

class ArrayMapFunction : public LambdaFunction {

ENABLE_FACTORY_CREATOR(ArrayMapFunction);

public:

~ArrayMapFunction() override = default;

static constexpr auto name = "array_map";

static LambdaFunctionPtr create() { return std::make_shared<ArrayMapFunction>(); }

std::string get_name() const override { return name; }

Status execute(VExprContext* context, const Block* block, const Selector* expr_selector,

size_t count, ColumnPtr& result_column, const DataTypePtr& result_type,

const VExprSPtrs& children) const override {

LambdaArgs args_info;

// collect used slot ref in lambda function body

std::vector<int>& output_slot_ref_indexs = args_info.output_slot_ref_indexs;

_collect_slot_ref_column_id(children[0], output_slot_ref_indexs);

int gap = 0;

if (!output_slot_ref_indexs.empty()) {

auto max_id = std::ranges::max_element(output_slot_ref_indexs);

gap = *max_id + 1;

_set_column_ref_column_id(children[0], gap);

}

std::vector<std::string> names(gap);

DataTypes data_types(gap);

for (int i = 0; i < gap; ++i) {

if (_contains_column_id(output_slot_ref_indexs, i)) {

names[i] = block->get_by_position(i).name;

data_types[i] = block->get_by_position(i).type;

} else {

// padding some mock data to hold the position, like call block#rows function need

names[i] = "temp";

data_types[i] = std::make_shared<DataTypeUInt8>();

}

}

///* array_map(lambda,arg1,arg2,.....) *///

//1. child[1:end]->execute(src_block)

ColumnsWithTypeAndName arguments(children.size() - 1);

for (int i = 1; i < children.size(); ++i) {

ColumnPtr column;

RETURN_IF_ERROR(

children[i]->execute_column(context, block, expr_selector, count, column));

arguments[i - 1].column = column;

arguments[i - 1].type = children[i]->execute_type(block);

arguments[i - 1].name = children[i]->expr_name();

}

// used for save column array outside null map

auto outside_null_map = ColumnUInt8::create(

arguments[0].column->convert_to_full_column_if_const()->size(), 0);

// offset column

MutableColumnPtr array_column_offset;

size_t nested_array_column_rows = 0;

ColumnPtr first_array_offsets = nullptr;

//2. get the result column from executed expr, and the needed is nested column of array

std::vector<ColumnPtr> lambda_datas(arguments.size());

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

const auto& array_column_type_name = arguments[i];

auto column_array = array_column_type_name.column->convert_to_full_column_if_const();

auto type_array = array_column_type_name.type;

if (type_array->is_nullable()) {

// get the nullmap of nullable column

// hold the null column instead of a reference 'cause `column_array` will be assigned and freed below.

DORIS_CHECK(is_column_nullable(*column_array));

auto column_array_nullmap =

assert_cast<const ColumnNullable&>(*column_array).get_null_map_column_ptr();

// get the array column from nullable column

column_array = assert_cast<const ColumnNullable*>(column_array.get())

->get_nested_column_ptr();

// get the nested type from nullable type

type_array = assert_cast<const DataTypeNullable*>(array_column_type_name.type.get())

->get_nested_type();

// need to union nullmap from all columns

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

column_array_nullmap->get_data());

}

// here is the array column

const auto& col_array = assert_cast<const ColumnArray&>(*column_array);

const auto& col_type = assert_cast<const DataTypeArray&>(*type_array);

if (i == 0) {

nested_array_column_rows = col_array.get_data_ptr()->size();

first_array_offsets = col_array.get_offsets_ptr();

const auto& off_data = col_array.get_offsets_column();

array_column_offset = off_data.clone_resized(col_array.get_offsets_column().size());

args_info.offsets_ptr = &col_array.get_offsets();

} else {

// select array_map((x,y)->x+y,c_array1,[0,1,2,3]) from array_test2;

// c_array1: [0,1,2,3,4,5,6,7,8,9]

const auto& array_offsets =

assert_cast<const ColumnArray::ColumnOffsets&>(*first_array_offsets)

.get_data();

if (nested_array_column_rows != col_array.get_data_ptr()->size() ||

(!array_offsets.empty() &&

memcmp(array_offsets.data(), col_array.get_offsets().data(),

sizeof(array_offsets[0]) * array_offsets.size()) != 0)) {

return Status::InvalidArgument(

"in array map function, the input column size "

"are "

"not equal completely, nested column data rows 1st size is {}, {}th "

"size is {}.",

nested_array_column_rows, i + 1, col_array.get_data_ptr()->size());

}

}

lambda_datas[i] = col_array.get_data_ptr();

names.push_back("R" + array_column_type_name.name);

data_types.push_back(col_type.get_nested_type());

}

// if column_array is NULL, we know the array_data_column will not write any data,

// so the column is empty. eg : (x) -> concat('|',x + "1"). if still execute the lambda function, will cause the bolck rows are not equal

// the x column is empty, but "|" is const literal, size of column is 1, so the block rows is 1, but the x column is empty, will be coredump.

if (std::any_of(lambda_datas.begin(), lambda_datas.end(),

[](const auto& v) { return v->empty(); })) {

DataTypePtr nested_type;

bool is_nullable = result_type->is_nullable();

if (is_nullable) {

nested_type =

assert_cast<const DataTypeNullable*>(result_type.get())->get_nested_type();

} else {

nested_type = result_type;

}

auto empty_nested_column = assert_cast<const DataTypeArray*>(nested_type.get())

->get_nested_type()

->create_column();

auto result_array_column = ColumnArray::create(std::move(empty_nested_column),

std::move(array_column_offset));

if (is_nullable) {

result_column = ColumnNullable::create(std::move(result_array_column),

std::move(outside_null_map));

} else {

result_column = std::move(result_array_column);

}

return Status::OK();

}

MutableColumnPtr result_col = nullptr;

DataTypePtr res_type;

//process first row

args_info.array_start = (*args_info.offsets_ptr)[args_info.current_row_idx - 1];

args_info.cur_size =

(*args_info.offsets_ptr)[args_info.current_row_idx] - args_info.array_start;

// lambda block to exectute the lambda, and reuse the memory

Block lambda_block;

auto column_size = names.size();

MutableColumns columns(column_size);

do {

bool mem_reuse = lambda_block.mem_reuse();

for (int i = 0; i < column_size; i++) {

if (mem_reuse) {

columns[i] = lambda_block.get_by_position(i).column->assert_mutable();

} else {

if (_contains_column_id(output_slot_ref_indexs, i) || i >= gap) {

// TODO: maybe could create const column, so not insert_many_from when extand data

// but now here handle batch_size of array nested data every time, so maybe have different rows

columns[i] = data_types[i]->create_column();

} else {

columns[i] = data_types[i]

->create_column_const_with_default_value(0)

->assert_mutable();

}

}

}

// batch_size of array nested data every time inorder to avoid memory overflow

while (columns[gap]->size() < batch_size) {

long max_step = batch_size - columns[gap]->size();

long current_step = std::min(

max_step, (long)(args_info.cur_size - args_info.current_offset_in_array));

size_t pos = args_info.array_start + args_info.current_offset_in_array;

for (int i = 0; i < arguments.size() && current_step > 0; ++i) {

columns[gap + i]->insert_range_from(*lambda_datas[i], pos, current_step);

}

args_info.current_offset_in_array += current_step;

args_info.current_repeat_times += current_step;

if (args_info.current_offset_in_array >= args_info.cur_size) {

args_info.current_row_eos = true;

}

_extend_data(columns, block, args_info.current_repeat_times, gap,

args_info.current_row_idx, output_slot_ref_indexs);

args_info.current_repeat_times = 0;

if (args_info.current_row_eos) {

//current row is end of array, move to next row

args_info.current_row_idx++;

args_info.current_offset_in_array = 0;

if (args_info.current_row_idx >= count) {

break;

}

args_info.current_row_eos = false;

args_info.array_start = (*args_info.offsets_ptr)[args_info.current_row_idx - 1];

args_info.cur_size = (*args_info.offsets_ptr)[args_info.current_row_idx] -

args_info.array_start;

}

}

if (!mem_reuse) {

for (int i = 0; i < column_size; ++i) {

lambda_block.insert(

ColumnWithTypeAndName(std::move(columns[i]), data_types[i], names[i]));

}

}

//3. child[0]->execute(new_block)

ColumnPtr res_col;

// lambda body executes on the internal lambda_block, not the original block.

// The outer expr_selector is irrelevant here, so pass nullptr.

RETURN_IF_ERROR(children[0]->execute_column(context, &lambda_block, nullptr,

lambda_block.rows(), res_col));

res_col = res_col->convert_to_full_column_if_const();

res_type = children[0]->execute_type(&lambda_block);

if (!result_col) {

result_col = res_col->clone_empty();

}

result_col->insert_range_from(*res_col, 0, res_col->size());

lambda_block.clear_column_data(column_size);

} while (args_info.current_row_idx < count);

//4. get the result column after execution, reassemble it into a new array column, and return.

if (result_type->is_nullable()) {

if (res_type->is_nullable()) {

result_column = ColumnNullable::create(

ColumnArray::create(std::move(result_col), std::move(array_column_offset)),

std::move(outside_null_map));

} else {

// deal with eg: select array_map(x -> x is null, [null, 1, 2]);

// need to create the nested column null map for column array

auto nested_null_map = ColumnUInt8::create(result_col->size(), 0);

result_column = ColumnNullable::create(

ColumnArray::create(ColumnNullable::create(std::move(result_col),

std::move(nested_null_map)),

std::move(array_column_offset)),

std::move(outside_null_map));

}

} else {

if (res_type->is_nullable()) {

result_column =

ColumnArray::create(std::move(result_col), std::move(array_column_offset));

} else {

auto nested_null_map = ColumnUInt8::create(result_col->size(), 0);

result_column = ColumnArray::create(

ColumnNullable::create(std::move(result_col), std::move(nested_null_map)),

std::move(array_column_offset));

}

}

return Status::OK();

}

private:

bool _contains_column_id(const std::vector<int>& output_slot_ref_indexs, int id) const {

const auto it = std::find(output_slot_ref_indexs.begin(), output_slot_ref_indexs.end(), id);

return it != output_slot_ref_indexs.end();

}

void _set_column_ref_column_id(VExprSPtr expr, int gap) const {

for (const auto& child : expr->children()) {

if (child->is_column_ref()) {

auto* ref = static_cast<VColumnRef*>(child.get());

ref->set_gap(gap);

} else {

_set_column_ref_column_id(child, gap);

}

}

}

void _collect_slot_ref_column_id(VExprSPtr expr,

std::vector<int>& output_slot_ref_indexs) const {

for (const auto& child : expr->children()) {

if (child->is_slot_ref()) {

const auto* ref = static_cast<VSlotRef*>(child.get());

output_slot_ref_indexs.push_back(ref->column_id());

} else {

_collect_slot_ref_column_id(child, output_slot_ref_indexs);

}

}

}

void _extend_data(std::vector<MutableColumnPtr>& columns, const Block* block,

int current_repeat_times, int size, int64_t current_row_idx,

const std::vector<int>& output_slot_ref_indexs) const {

if (!current_repeat_times || !size) {

return;

}

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

if (_contains_column_id(output_slot_ref_indexs, i)) {

auto src_column =

block->get_by_position(i).column->convert_to_full_column_if_const();

columns[i]->insert_many_from(*src_column, current_row_idx, current_repeat_times);

} else {

// must be column const

DCHECK(is_column_const(*columns[i]));

columns[i]->resize(columns[i]->size() + current_repeat_times);

}

}

}

};

void register_function_array_map(doris::LambdaFunctionFactory& factory) {

factory.register_function<ArrayMapFunction>();

}

} // namespace doris