{

const ColumnPtr & key_column;

const ColumnPtr & val_column;

const IColumn::Offsets & key_offsets;

const IColumn::Offsets & val_offsets;

bool is_const;

{

ADD = 0,

SUBTRACT = 1

{

static FunctionPtr create(ContextPtr, OpTypes op_type_) { return std::make_shared<FunctionMapOp>(op_type_); }

explicit FunctionMapOp(OpTypes op_type_) : op_type(op_type_) {}

const OpTypes op_type;

String getName() const override { return (op_type == OpTypes::ADD) ? "mapAdd" : "mapSubtract"; }

size_t getNumberOfArguments() const override { return 0; }

bool isVariadic() const override { return true; }

bool useDefaultImplementationForConstants() const override { return true; }

bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }

void checkTypes(

DataTypePtr & key_type, DataTypePtr & promoted_val_type, const DataTypePtr & check_key_type, DataTypePtr & check_val_type) const

{

if (!(check_key_type->equals(*key_type)))

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Expected same {} type for all keys in {}",

key_type->getName(), getName());

WhichDataType which_val(promoted_val_type);

WhichDataType which_ch_val(check_val_type);

if (which_ch_val.isFloat() != which_val.isFloat())

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "All value types in {} should be either or float or integer",

getName());

if (!(check_val_type->equals(*promoted_val_type)))

{

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "All value types in {} should be promotable to {}, got {}",

getName(), promoted_val_type->getName(), check_val_type->getName());

}

}

DataTypePtr getReturnTypeForTuples(const DataTypes & arguments) const

{

DataTypePtr key_type;

DataTypePtr val_type;

DataTypePtr res;

for (const auto & arg : arguments)

{

const DataTypeArray * k = nullptr;

const DataTypeArray * v = nullptr;

const DataTypeTuple * tup = checkAndGetDataType<DataTypeTuple>(arg.get());

if (!tup)

throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "{} accepts at least two map tuples", getName());

auto elems = tup->getElements();

if (elems.size() != 2)

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Each tuple in {} arguments should consist of two arrays",

getName());

k = checkAndGetDataType<DataTypeArray>(elems[0].get());

v = checkAndGetDataType<DataTypeArray>(elems[1].get());

if (!k || !v)

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Each tuple in {} arguments should consist of two arrays",

getName());

const auto & result_type = v->getNestedType();

if (!result_type->canBePromoted())

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Values to be summed are expected to be Numeric, Float or Decimal.");

auto promoted_val_type = result_type->promoteNumericType();

if (!key_type)

{

key_type = k->getNestedType();

val_type = promoted_val_type;

res = std::make_shared<DataTypeTuple>(

DataTypes{std::make_shared<DataTypeArray>(k->getNestedType()), std::make_shared<DataTypeArray>(promoted_val_type)});

}

else

checkTypes(key_type, val_type, k->getNestedType(), promoted_val_type);

}

return res;

}

DataTypePtr getReturnTypeForMaps(const DataTypes & arguments) const

{

DataTypePtr key_type;

DataTypePtr val_type;

DataTypePtr res;

for (const auto & arg : arguments)

{

const auto * map = checkAndGetDataType<DataTypeMap>(arg.get());

if (!map)

throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "{} accepts at least two maps", getName());

const auto & v = map->getValueType();

if (!v->canBePromoted())

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Values to be summed are expected to be Numeric, Float or Decimal.");

auto promoted_val_type = v->promoteNumericType();

if (!key_type)

{

key_type = map->getKeyType();

val_type = promoted_val_type;

res = std::make_shared<DataTypeMap>(DataTypes({key_type, promoted_val_type}));

}

else

checkTypes(key_type, val_type, map->getKeyType(), promoted_val_type);

}

return res;

}

DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override

{

if (arguments.size() < 2)

throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "{} accepts at least two maps or map tuples", getName());

if (arguments[0]->getTypeId() == TypeIndex::Tuple)

return getReturnTypeForTuples(arguments);

if (arguments[0]->getTypeId() == TypeIndex::Map)

return getReturnTypeForMaps(arguments);

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "{} only accepts maps", getName());

}

template <typename KeyType, typename ValType>

ColumnPtr execute2(size_t row_count, TupleMaps & args, const DataTypePtr res_type) const

{

MutableColumnPtr res_column = res_type->createColumn();

IColumn *to_keys_data = nullptr;

IColumn *to_vals_data = nullptr;

ColumnArray::Offsets * to_keys_offset = nullptr;

ColumnArray::Offsets * to_vals_offset = nullptr;

// prepare output destinations

if (res_type->getTypeId() == TypeIndex::Tuple)

{

auto * to_tuple = assert_cast<ColumnTuple *>(res_column.get());

auto & to_keys_arr = assert_cast<ColumnArray &>(to_tuple->getColumn(0));

to_keys_data = &to_keys_arr.getData();

to_keys_offset = &to_keys_arr.getOffsets();

auto & to_vals_arr = assert_cast<ColumnArray &>(to_tuple->getColumn(1));

to_vals_data = &to_vals_arr.getData();

to_vals_offset = &to_vals_arr.getOffsets();

}

else

{

chassert(res_type->getTypeId() == TypeIndex::Map);

auto * to_map = assert_cast<ColumnMap *>(res_column.get());

auto & to_wrapper_arr = to_map->getNestedColumn();

to_keys_offset = &to_wrapper_arr.getOffsets();

auto & to_map_tuple = to_map->getNestedData();

to_keys_data = &to_map_tuple.getColumn(0);

to_vals_data = &to_map_tuple.getColumn(1);

}

MapWithMemoryTracking<KeyType, ValType> summing_map;

for (size_t i = 0; i < row_count; ++i)

{

[[maybe_unused]] bool first = true;

for (auto & arg : args)

{

size_t offset = 0;

size_t len = arg.key_offsets[0];

if (!arg.is_const)

{

offset = arg.key_offsets[i - 1];

len = arg.key_offsets[i] - offset;

if (arg.val_offsets[i] != arg.key_offsets[i])

throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Key and value array should have same amount of elements");

}

Field temp_val;

for (size_t j = 0; j < len; ++j)

{

KeyType key{};

if constexpr (std::is_same_v<KeyType, String>)

{

if (const auto * col_fixed = checkAndGetColumn<ColumnFixedString>(arg.key_column.get()))

key = col_fixed->getDataAt(offset + j);

else if (const auto * col_str = checkAndGetColumn<ColumnString>(arg.key_column.get()))

key = col_str->getDataAt(offset + j);

else // should not happen

throw Exception(ErrorCodes::LOGICAL_ERROR,

"Expected String or FixedString, got {} in {}",

arg.key_column->getDataType(), getName());

}

else

{

key = assert_cast<const ColumnVector<KeyType> *>(arg.key_column.get())->getData()[offset + j];

}

arg.val_column->get(offset + j, temp_val);

ValType value = temp_val.safeGet<ValType>();

if (op_type == OpTypes::ADD)

{

const auto [it, inserted] = summing_map.insert({key, value});

if (!inserted)

it->second = common::addIgnoreOverflow(it->second, value);

}

else

{

const auto [it, inserted] = summing_map.insert({key, first ? value : common::negateIgnoreOverflow(value)});

if (!inserted)

{

if (first)

it->second = common::addIgnoreOverflow(it->second, value);

else

it->second = common::subIgnoreOverflow(it->second, value);

}

}

}

first = false;

}

for (const auto & elem : summing_map)

{

to_keys_data->insert(elem.first);

to_vals_data->insert(elem.second);

}

to_keys_offset->push_back(to_keys_data->size());

summing_map.clear();

}

if (to_vals_offset)

{

// same offsets as in keys

to_vals_offset->insert(to_keys_offset->begin(), to_keys_offset->end());

}

return res_column;

}

template <typename KeyType>

ColumnPtr execute1(size_t row_count, const DataTypePtr res_type, const DataTypePtr res_value_type, TupleMaps & args) const

{

switch (res_value_type->getTypeId())

{

case TypeIndex::Int64:

return execute2<KeyType, Int64>(row_count, args, res_type);

case TypeIndex::Int128:

return execute2<KeyType, Int128>(row_count, args, res_type);

case TypeIndex::Int256:

return execute2<KeyType, Int256>(row_count, args, res_type);

case TypeIndex::UInt64:

return execute2<KeyType, UInt64>(row_count, args, res_type);

case TypeIndex::UInt128:

return execute2<KeyType, UInt128>(row_count, args, res_type);

case TypeIndex::UInt256:

return execute2<KeyType, UInt256>(row_count, args, res_type);

case TypeIndex::Float64:

return execute2<KeyType, Float64>(row_count, args, res_type);

case TypeIndex::Decimal128:

return execute2<KeyType, Decimal128>(row_count, args, res_type);

case TypeIndex::Decimal256:

return execute2<KeyType, Decimal256>(row_count, args, res_type);

default:

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal column type {} for values in arguments of function {}",

res_value_type->getName(), getName());

}

}

ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const override

{

DataTypePtr key_type;

size_t row_count = 0;

const DataTypeTuple * tup_type = checkAndGetDataType<DataTypeTuple>((arguments[0]).type.get());

DataTypePtr res_type;

DataTypePtr res_value_type;

TupleMaps args{};

args.reserve(arguments.size());

//prepare columns, extract data columns for direct access and put them to the vector

if (tup_type)

{

const DataTypeArray * key_array_type = checkAndGetDataType<DataTypeArray>(tup_type->getElements()[0].get());

const DataTypeArray * val_array_type = checkAndGetDataType<DataTypeArray>(tup_type->getElements()[1].get());

/* determine output type */

res_value_type = val_array_type->getNestedType()->promoteNumericType();

res_type = std::make_shared<DataTypeTuple>(DataTypes{

std::make_shared<DataTypeArray>(key_array_type->getNestedType()), std::make_shared<DataTypeArray>(res_value_type)});

for (const auto & col : arguments)

{

const ColumnTuple * tup = nullptr;

bool is_const = isColumnConst(*col.column);

if (is_const)

{

const auto * c = assert_cast<const ColumnConst *>(col.column.get());

tup = assert_cast<const ColumnTuple *>(c->getDataColumnPtr().get());

}

else

tup = assert_cast<const ColumnTuple *>(col.column.get());

const auto & arr1 = assert_cast<const ColumnArray &>(tup->getColumn(0));

const auto & arr2 = assert_cast<const ColumnArray &>(tup->getColumn(1));

const auto & key_offsets = arr1.getOffsets();

const auto & key_column = arr1.getDataPtr();

const auto & val_offsets = arr2.getOffsets();

const auto & val_column = arr2.getDataPtr();

args.push_back({key_column, val_column, key_offsets, val_offsets, is_const});

}

key_type = key_array_type->getNestedType();

}

else

{

const DataTypeMap * map_type = checkAndGetDataType<DataTypeMap>((arguments[0]).type.get());

if (map_type)

{

key_type = map_type->getKeyType();

res_value_type = map_type->getValueType()->promoteNumericType();

res_type = std::make_shared<DataTypeMap>(DataTypes{map_type->getKeyType(), res_value_type});

for (const auto & col : arguments)

{

const ColumnMap * map = nullptr;

bool is_const = isColumnConst(*col.column);

if (is_const)

{

const auto * c = assert_cast<const ColumnConst *>(col.column.get());

map = assert_cast<const ColumnMap *>(c->getDataColumnPtr().get());

}

else

map = assert_cast<const ColumnMap *>(col.column.get());

const auto & map_arr = map->getNestedColumn();

const auto & key_offsets = map_arr.getOffsets();

const auto & val_offsets = key_offsets;

const auto & map_tup = map->getNestedData();

const auto & key_column = map_tup.getColumnPtr(0);

const auto & val_column = map_tup.getColumnPtr(1);

args.push_back({key_column, val_column, key_offsets, val_offsets, is_const});

}

}

else

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal column type {} in arguments of function {}",

arguments[0].type->getName(), getName());

}

// we can check const columns before any processing

for (auto & arg : args)

{

if (arg.is_const)

{

if (arg.val_offsets[0] != arg.key_offsets[0])

throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Key and value array should have same amount of elements");

}

}

row_count = arguments[0].column->size();

switch (key_type->getTypeId())

{

case TypeIndex::Enum8:

case TypeIndex::Int8:

return execute1<Int8>(row_count, res_type, res_value_type, args);

case TypeIndex::Enum16:

case TypeIndex::Int16:

return execute1<Int16>(row_count, res_type, res_value_type, args);

case TypeIndex::Int32:

return execute1<Int32>(row_count, res_type, res_value_type, args);

case TypeIndex::Int64:

return execute1<Int64>(row_count, res_type, res_value_type, args);

case TypeIndex::Int128:

return execute1<Int128>(row_count, res_type, res_value_type, args);

case TypeIndex::Int256:

return execute1<Int256>(row_count, res_type, res_value_type, args);

case TypeIndex::UInt8:

return execute1<UInt8>(row_count, res_type, res_value_type, args);

case TypeIndex::Date:

case TypeIndex::UInt16:

return execute1<UInt16>(row_count, res_type, res_value_type, args);

case TypeIndex::DateTime:

case TypeIndex::UInt32:

return execute1<UInt32>(row_count, res_type, res_value_type, args);

case TypeIndex::UInt64:

return execute1<UInt64>(row_count, res_type, res_value_type, args);

case TypeIndex::UInt128:

return execute1<UInt128>(row_count, res_type, res_value_type, args);

case TypeIndex::UInt256:

return execute1<UInt256>(row_count, res_type, res_value_type, args);

case TypeIndex::UUID:

return execute1<UUID>(row_count, res_type, res_value_type, args);

case TypeIndex::FixedString:

case TypeIndex::String:

return execute1<String>(row_count, res_type, res_value_type, args);

default:

throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal column type {} for keys in arguments of function {}",

key_type->getName(), getName());

}

}

{

/// mapAdd function documentation

FunctionDocumentation::Description description_mapAdd = R"(

)";

FunctionDocumentation::Syntax syntax_mapAdd = "mapAdd(arg1[, arg2, ...])";

FunctionDocumentation::Arguments arguments_mapAdd = {

{"arg1[, arg2, ...]", "Maps or tuples of two arrays in which items in the first array represent keys, and the second array contains values for each key.", {"Map(K, V)", "Tuple(Array(T), Array(T))"}}

};

FunctionDocumentation::ReturnedValue returned_value_mapAdd = {"Returns a map or returns a tuple, where the first array contains the sorted keys and the second array contains values.", {"Map(K, V)", "Tuple(Array(T), Array(T))"}};

FunctionDocumentation::Examples examples_mapAdd = {

{"With Map type", "SELECT mapAdd(map(1, 1), map(1, 1))", "{1:2}"},

{"With tuple", "SELECT mapAdd(([toUInt8(1), 2], [1, 1]), ([toUInt8(1), 2], [1, 1]))", "([1,2],[2,2])"}

};

FunctionDocumentation::IntroducedIn introduced_in_mapAdd = {20, 7};

FunctionDocumentation::Category category_mapAdd = FunctionDocumentation::Category::Map;

FunctionDocumentation documentation_mapAdd = {description_mapAdd, syntax_mapAdd, arguments_mapAdd, {}, returned_value_mapAdd, examples_mapAdd, introduced_in_mapAdd, category_mapAdd};

factory.registerFunction("mapAdd", [](ContextPtr context){ return FunctionMapOp::create(context, OpTypes::ADD); }, documentation_mapAdd);

/// mapSubtract function documentation

FunctionDocumentation::Description description_mapSubtract = R"(

)";

FunctionDocumentation::Syntax syntax_mapSubtract = "mapSubtract(arg1[, arg2, ...])";

FunctionDocumentation::Arguments arguments_mapSubtract = {

{"arg1[, arg2, ...]", "Maps or tuples of two arrays in which items in the first array represent keys, and the second array contains values for each key.", {"Map(K, V)", "Tuple(Array(T), Array(T))"}}

};

FunctionDocumentation::ReturnedValue returned_value_mapSubtract = {"Returns one map or tuple, where the first array contains the sorted keys and the second array contains values.", {"Map(K, V)", "Tuple(Array(T), Array(T))"}};

FunctionDocumentation::Examples examples_mapSubtract = {

{"With Map type", "SELECT mapSubtract(map(1, 1), map(1, 1))", "{1:0}"},

{"With tuple map", "SELECT mapSubtract(([toUInt8(1), 2], [toInt32(1), 1]), ([toUInt8(1), 2], [toInt32(2), 1]))", "([1,2],[-1,0])"}

};

FunctionDocumentation::IntroducedIn introduced_in_mapSubtract = {20, 7};

FunctionDocumentation::Category category_mapSubtract = FunctionDocumentation::Category::Map;

FunctionDocumentation documentation_mapSubtract = {description_mapSubtract, syntax_mapSubtract, arguments_mapSubtract, {}, returned_value_mapSubtract, examples_mapSubtract, introduced_in_mapSubtract, category_mapSubtract};

factory.registerFunction("mapSubtract", [](ContextPtr context){ return FunctionMapOp::create(context, OpTypes::SUBTRACT); }, documentation_mapSubtract);