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

#include <algorithm>

#include <functional>

#include <memory>

#include <type_traits>

#include <utility>

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

#include "core/column/column_complex.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_bitmap.h"

#include "core/data_type/data_type_nullable.h"

#include "core/data_type/data_type_number.h"

#include "core/types.h"

#include "core/value/bitmap_value.h"

#include "exprs/aggregate/aggregate_function.h"

#include "exprs/function/function.h"

#include "exprs/function/simple_function_factory.h"

namespace doris {

class FunctionContext;

} // namespace doris

namespace doris {

// currently only bitmap_or and bitmap_or_count will call this function,

// other bitmap functions will use default implementation for nulls

#define BITMAP_OR_NULLABLE(nullable, input_rows_count, res, op) \

const auto& nested_col_ptr = nullable->get_nested_column_ptr(); \

const auto* __restrict null_map_data = nullable->get_null_map_data().data(); \

const auto& mid_data = assert_cast<const ColumnBitmap*>(nested_col_ptr.get())->get_data(); \

for (size_t row = 0; row < input_rows_count; ++row) { \

if (!null_map_data[row]) { \

res[row] op mid_data[row]; \

} \

}

#define BITMAP_FUNCTION_VARIADIC(CLASS, FUNCTION_NAME, OP) \

struct CLASS { \

static constexpr auto name = #FUNCTION_NAME; \

using ResultDataType = DataTypeBitMap; \

static Status vector_vector(ColumnPtr argument_columns[], size_t col_size, \

size_t input_rows_count, std::vector<BitmapValue>& res, \

IColumn* res_nulls) { \

std::vector<const ColumnUInt8::value_type*> null_map_datas(col_size); \

int nullable_cols_count = 0; \

ColumnUInt8::value_type* __restrict res_nulls_data = nullptr; \

if (res_nulls) { \

res_nulls_data = assert_cast<ColumnUInt8*>(res_nulls)->get_data().data(); \

} \

if (auto* nullable = check_and_get_column<ColumnNullable>(*argument_columns[0])) { \

null_map_datas[nullable_cols_count++] = nullable->get_null_map_data().data(); \

BITMAP_OR_NULLABLE(nullable, input_rows_count, res, =); \

} else { \

const auto& mid_data = \

assert_cast<const ColumnBitmap*>(argument_columns[0].get())->get_data(); \

for (size_t row = 0; row < input_rows_count; ++row) { \

res[row] = mid_data[row]; \

} \

} \

for (size_t col = 1; col < col_size; ++col) { \

if (auto* nullable = \

check_and_get_column<ColumnNullable>(*argument_columns[col])) { \

null_map_datas[nullable_cols_count++] = nullable->get_null_map_data().data(); \

BITMAP_OR_NULLABLE(nullable, input_rows_count, res, OP); \

} else { \

const auto& col_data = \

assert_cast<const ColumnBitmap*>(argument_columns[col].get()) \

->get_data(); \

for (size_t row = 0; row < input_rows_count; ++row) { \

res[row] OP col_data[row]; \

} \

} \

} \

if (res_nulls_data && nullable_cols_count == col_size) { \

const auto* null_map_data = null_map_datas[0]; \

for (size_t row = 0; row < input_rows_count; ++row) { \

res_nulls_data[row] = null_map_data[row]; \

} \

for (int i = 1; i < nullable_cols_count; ++i) { \

null_map_data = null_map_datas[i]; \

for (size_t row = 0; row < input_rows_count; ++row) { \

res_nulls_data[row] &= null_map_data[row]; \

} \

} \

} \

return Status::OK(); \

} \

}

#define BITMAP_FUNCTION_COUNT_VARIADIC(CLASS, FUNCTION_NAME, OP) \

struct CLASS { \

static constexpr auto name = #FUNCTION_NAME; \

using ResultDataType = DataTypeInt64; \

using TData = std::vector<BitmapValue>; \

using ResTData = typename ColumnInt64::Container; \

static Status vector_vector(ColumnPtr argument_columns[], size_t col_size, \

size_t input_rows_count, ResTData& res, IColumn* res_nulls) { \

TData vals; \

if (auto* nullable = check_and_get_column<ColumnNullable>(*argument_columns[0])) { \

vals.resize(input_rows_count); \

BITMAP_OR_NULLABLE(nullable, input_rows_count, vals, =); \

} else { \

vals = assert_cast<const ColumnBitmap*>(argument_columns[0].get())->get_data(); \

} \

for (size_t col = 1; col < col_size; ++col) { \

if (auto* nullable = \

check_and_get_column<ColumnNullable>(*argument_columns[col])) { \

BITMAP_OR_NULLABLE(nullable, input_rows_count, vals, OP); \

} else { \

const auto& col_data = \

assert_cast<const ColumnBitmap*>(argument_columns[col].get()) \

->get_data(); \

for (size_t row = 0; row < input_rows_count; ++row) { \

vals[row] OP col_data[row]; \

} \

} \

} \

for (size_t row = 0; row < input_rows_count; ++row) { \

res[row] = vals[row].cardinality(); \

} \

return Status::OK(); \

} \

}

BITMAP_FUNCTION_VARIADIC(BitmapOr, bitmap_or, |=);

BITMAP_FUNCTION_VARIADIC(BitmapAnd, bitmap_and, &=);

BITMAP_FUNCTION_VARIADIC(BitmapXor, bitmap_xor, ^=);

BITMAP_FUNCTION_COUNT_VARIADIC(BitmapOrCount, bitmap_or_count, |=);

BITMAP_FUNCTION_COUNT_VARIADIC(BitmapAndCount, bitmap_and_count, &=);

BITMAP_FUNCTION_COUNT_VARIADIC(BitmapXorCount, bitmap_xor_count, ^=);

Status execute_bitmap_op_count_null_to_zero(

FunctionContext* context, Block& block, const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count,

const std::function<Status(FunctionContext*, Block&, const ColumnNumbers&, size_t, size_t)>&

exec_impl_func);

template <typename Impl>

class FunctionBitMapVariadic : public IFunction {

public:

static constexpr auto name = Impl::name;

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

String get_name() const override { return name; }

size_t get_number_of_arguments() const override { return 0; }

bool is_variadic() const override { return true; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

using ResultDataType = typename Impl::ResultDataType;

if (std::is_same_v<Impl, BitmapOr> || is_count()) {

bool return_nullable = false;

// result is nullable only when any columns is nullable for bitmap_or and bitmap_or_count

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

if (arguments[i]->is_nullable()) {

return_nullable = true;

break;

}

}

auto result_type = std::make_shared<ResultDataType>();

return return_nullable ? make_nullable(result_type) : result_type;

} else {

return std::make_shared<ResultDataType>();

}

}

bool use_default_implementation_for_nulls() const override {

// result is null only when all columns is null for bitmap_or.

// for count functions, result is always not null, and if the bitmap op result is null,

// the count is 0

return !static_cast<bool>(std::is_same_v<Impl, BitmapOr> || is_count());

}

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

uint32_t result, size_t input_rows_count) const override {

if (std::is_same_v<Impl, BitmapAndCount> || std::is_same_v<Impl, BitmapXorCount>) {

auto impl_func = [&](FunctionContext* context, Block& block,

const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) {

return execute_impl_internal(context, block, arguments, result, input_rows_count);

};

return execute_bitmap_op_count_null_to_zero(context, block, arguments, result,

input_rows_count, impl_func);

} else {

return execute_impl_internal(context, block, arguments, result, input_rows_count);

}

}

Status execute_impl_internal(FunctionContext* context, Block& block,

const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) const {

size_t argument_size = arguments.size();

std::vector<ColumnPtr> argument_columns(argument_size);

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

argument_columns[i] =

block.get_by_position(arguments[i]).column->convert_to_full_column_if_const();

}

using ResultDataType = typename Impl::ResultDataType; //DataTypeBitMap or DataTypeInt64

using ColVecResult = std::conditional_t<is_complex_v<ResultDataType::PType>,

ColumnComplexType<ResultDataType::PType>,

ColumnVector<ResultDataType::PType>>;

typename ColVecResult::MutablePtr col_res = nullptr;

typename ColumnUInt8::MutablePtr col_res_nulls;

auto& result_info = block.get_by_position(result);

// special case for bitmap_or and bitmap_or_count

if (!use_default_implementation_for_nulls() && result_info.type->is_nullable()) {

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

}

col_res = ColVecResult::create();

auto& vec_res = col_res->get_data();

vec_res.resize(input_rows_count);

RETURN_IF_ERROR(Impl::vector_vector(argument_columns.data(), argument_size,

input_rows_count, vec_res, col_res_nulls.get()));

if (!use_default_implementation_for_nulls() && result_info.type->is_nullable()) {

block.replace_by_position(

result, ColumnNullable::create(std::move(col_res), std::move(col_res_nulls)));

} else {

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

}

return Status::OK();

}

private:

bool is_count() const {

return (std::is_same_v<Impl, BitmapOrCount> || std::is_same_v<Impl, BitmapAndCount> ||

std::is_same_v<Impl, BitmapXorCount>);

}

};

using FunctionBitmapOr = FunctionBitMapVariadic<BitmapOr>;

using FunctionBitmapXor = FunctionBitMapVariadic<BitmapXor>;

using FunctionBitmapAnd = FunctionBitMapVariadic<BitmapAnd>;

using FunctionBitmapOrCount = FunctionBitMapVariadic<BitmapOrCount>;

using FunctionBitmapAndCount = FunctionBitMapVariadic<BitmapAndCount>;

using FunctionBitmapXorCount = FunctionBitMapVariadic<BitmapXorCount>;

void register_function_bitmap_variadic(SimpleFunctionFactory& factory) {

factory.register_function<FunctionBitmapOr>();

factory.register_function<FunctionBitmapXor>();

factory.register_function<FunctionBitmapAnd>();

factory.register_function<FunctionBitmapOrCount>();

factory.register_function<FunctionBitmapAndCount>();

factory.register_function<FunctionBitmapXorCount>();

}

} // namespace doris