FunctionUnaryArithmetic.h

#pragma once

#include <Columns/ColumnDecimal.h>
#include <Columns/ColumnFixedString.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnVector.h>
#include <Columns/ColumnsNumber.h>
#include <DataTypes/DataTypeFixedString.h>
#include <DataTypes/DataTypeInterval.h>
#include <DataTypes/DataTypeString.h>
#include <DataTypes/DataTypesDecimal.h>
#include <DataTypes/DataTypesNumber.h>
#include <DataTypes/Native.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/IFunction.h>
#include <Functions/IsOperation.h>
#include <Functions/castTypeToEither.h>

#include "config.h"
#include <Common/TargetSpecific.h>

#if USE_EMBEDDED_COMPILER
#    include <llvm/IR/IRBuilder.h>
#endif


namespace DB
{

namespace ErrorCodes
{
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
    extern const int LOGICAL_ERROR;
}

template <typename A, typename Op>
struct UnaryOperationImpl
{
    using ResultType = typename Op::ResultType;
    using ColVecA = ColumnVectorOrDecimal<A>;
    using ColVecC = ColumnVectorOrDecimal<ResultType>;
    using ArrayA = typename ColVecA::Container;
    using ArrayC = typename ColVecC::Container;

    MULTITARGET_FUNCTION_AVX512BW_AVX512F_AVX2_SSE42(
    MULTITARGET_FUNCTION_HEADER(static void NO_INLINE), vectorImpl, MULTITARGET_FUNCTION_BODY((const ArrayA & a, ArrayC & c) /// NOLINT
    {
        size_t size = a.size();
        for (size_t i = 0; i < size; ++i)
            c[i] = Op::apply(a[i]);
    }))

    static void NO_INLINE vector(const ArrayA & a, ArrayC & c)
    {
#if USE_MULTITARGET_CODE
        if (isArchSupported(TargetArch::AVX512BW))
        {
            vectorImplAVX512BW(a, c);
            return;
        }

        if (isArchSupported(TargetArch::AVX512F))
        {
            vectorImplAVX512F(a, c);
            return;
        }

        if (isArchSupported(TargetArch::AVX2))
        {
            vectorImplAVX2(a, c);
            return;
        }

        if (isArchSupported(TargetArch::SSE42))
        {
            vectorImplSSE42(a, c);
            return;
        }
#endif

        vectorImpl(a, c);
    }

    static void constant(A a, ResultType & c)
    {
        c = Op::apply(a);
    }
};


template <typename Op>
struct FixedStringUnaryOperationImpl
{
    MULTITARGET_FUNCTION_AVX512BW_AVX512F_AVX2_SSE42(
    MULTITARGET_FUNCTION_HEADER(static void NO_INLINE), vectorImpl, MULTITARGET_FUNCTION_BODY((const ColumnFixedString::Chars & a, /// NOLINT
        ColumnFixedString::Chars & c)
    {
        size_t size = a.size();
        for (size_t i = 0; i < size; ++i)
            c[i] = Op::apply(a[i]);
    }))

    static void NO_INLINE vector(const ColumnFixedString::Chars & a, ColumnFixedString::Chars & c)
    {
#if USE_MULTITARGET_CODE
        if (isArchSupported(TargetArch::AVX512BW))
        {
            vectorImplAVX512BW(a, c);
            return;
        }

        if (isArchSupported(TargetArch::AVX512F))
        {
            vectorImplAVX512F(a, c);
            return;
        }

        if (isArchSupported(TargetArch::AVX2))
        {
            vectorImplAVX2(a, c);
            return;
        }

        if (isArchSupported(TargetArch::SSE42))
        {
            vectorImplSSE42(a, c);
            return;
        }
#endif

        vectorImpl(a, c);
    }
};

template <typename Op>
struct StringUnaryOperationReduceImpl
{
    MULTITARGET_FUNCTION_AVX512BW_AVX512F_AVX2_SSE42(
        MULTITARGET_FUNCTION_HEADER(static UInt64 NO_INLINE),
        vectorImpl,
        MULTITARGET_FUNCTION_BODY((const UInt8 * start, const UInt8 * end) /// NOLINT
        {
            UInt64 res = 0;
            while (start < end)
                res += Op::apply(*start++);
            return res;
        }))

    static UInt64 NO_INLINE vector(const UInt8 * start, const UInt8 * end)
    {
#if USE_MULTITARGET_CODE
        if (isArchSupported(TargetArch::AVX512BW))
        {
            return vectorImplAVX512BW(start, end);
        }

        if (isArchSupported(TargetArch::AVX512F))
        {
            return vectorImplAVX512F(start, end);
        }

        if (isArchSupported(TargetArch::AVX2))
        {
            return vectorImplAVX2(start, end);
        }

        if (isArchSupported(TargetArch::SSE42))
        {
            return vectorImplSSE42(start, end);
        }
#endif

        return vectorImpl(start, end);
    }
};

template <typename FunctionName>
struct FunctionUnaryArithmeticMonotonicity;

/// Used to indicate undefined operation
struct InvalidType;


template <template <typename> class Op, typename Name, bool is_injective>
class FunctionUnaryArithmetic : public IFunction
{
    static constexpr bool allow_decimal = IsUnaryOperation<Op>::negate || IsUnaryOperation<Op>::abs || IsUnaryOperation<Op>::sign;
    static constexpr bool allow_string_or_fixed_string = Op<UInt8>::allow_string_or_fixed_string;
    static constexpr bool is_bit_count = IsUnaryOperation<Op>::bit_count;
    static constexpr bool is_sign_function = IsUnaryOperation<Op>::sign;

    ContextPtr context;

    template <typename F>
    static bool castType(const IDataType * type, F && f)
    {
        return castTypeToEither<
            DataTypeUInt8,
            DataTypeUInt16,
            DataTypeUInt32,
            DataTypeUInt64,
            DataTypeUInt128,
            DataTypeUInt256,
            DataTypeInt8,
            DataTypeInt16,
            DataTypeInt32,
            DataTypeInt64,
            DataTypeInt128,
            DataTypeInt256,
            DataTypeFloat32,
            DataTypeFloat64,
            DataTypeDecimal<Decimal32>,
            DataTypeDecimal<Decimal64>,
            DataTypeDecimal<Decimal128>,
            DataTypeDecimal<Decimal256>,
            DataTypeFixedString,
            DataTypeString,
            DataTypeInterval>(type, std::forward<F>(f));
    }

    static FunctionOverloadResolverPtr
    getFunctionForTupleArithmetic(const DataTypePtr & type, ContextPtr context)
    {
        if (!isTuple(type))
            return {};

        /// Special case when the function is negate, argument is tuple.
        /// We construct another function (example: tupleNegate) and call it.

        if constexpr (!IsUnaryOperation<Op>::negate)
            return {};

        return FunctionFactory::instance().get("tupleNegate", context);
    }

public:
    static constexpr auto name = Name::name;
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionUnaryArithmetic>(); }

    FunctionUnaryArithmetic() = default;

    explicit FunctionUnaryArithmetic(ContextPtr context_) : context(context_) {}

    String getName() const override
    {
        return name;
    }

    size_t getNumberOfArguments() const override { return 1; }
    bool isInjective(const ColumnsWithTypeAndName &) const override { return is_injective; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override
    {
        return false;
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        return getReturnTypeImplStatic(arguments, context);
    }

    static DataTypePtr getReturnTypeImplStatic(const DataTypes & arguments, ContextPtr context)
    {
        /// Special case when the function is negate, argument is tuple.
        if (auto function_builder = getFunctionForTupleArithmetic(arguments[0], context))
        {
            ColumnsWithTypeAndName new_arguments(1);

            new_arguments[0].type = arguments[0];

            auto function = function_builder->build(new_arguments);
            return function->getResultType();
        }

        DataTypePtr result;
        bool valid = castType(arguments[0].get(), [&](const auto & type)
        {
            using DataType = std::decay_t<decltype(type)>;
            if constexpr (std::is_same_v<DataTypeFixedString, DataType>)
            {
                if constexpr (!allow_string_or_fixed_string)
                    return false;
                /// For `bitCount`, when argument is FixedString, it's return type
                /// should be integer instead of FixedString, the return value is
                /// the sum of `bitCount` apply to each chars.
                else
                {
                    /// UInt16 can save bitCount of FixedString less than 8192,
                    /// it's should enough for almost all cases, and the setting
                    /// `allow_suspicious_fixed_string_types` is disabled by default.
                    if constexpr (is_bit_count)
                        result = std::make_shared<DataTypeUInt16>();
                    else
                        result = std::make_shared<DataType>(type.getN());
                }
            }
            else if constexpr (std::is_same_v<DataTypeString, DataType>)
            {
                if constexpr (!allow_string_or_fixed_string)
                    return false;
                else
                {
                    if constexpr (is_bit_count)
                        result = std::make_shared<DataTypeUInt64>();
                    else
                        result = std::make_shared<DataType>();
                }
            }
            else if constexpr (std::is_same_v<DataTypeInterval, DataType>)
            {
                if constexpr (!IsUnaryOperation<Op>::negate)
                    return false;
                result = std::make_shared<DataTypeInterval>(type.getKind());
            }
            else
            {
                using T0 = typename DataType::FieldType;

                if constexpr (IsDataTypeDecimal<DataType> && !is_sign_function)
                {
                    if constexpr (!allow_decimal)
                        return false;
                    result = std::make_shared<DataType>(type.getPrecision(), type.getScale());
                }
                else
                    result = std::make_shared<DataTypeNumber<typename Op<T0>::ResultType>>();
            }
            return true;
        });
        if (!valid)
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal type {} of argument of function {}",
                arguments[0]->getName(), String(name));
        return result;
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
    {
        /// Special case when the function is negate, argument is tuple.
        if (auto function_builder = getFunctionForTupleArithmetic(arguments[0].type, context))
        {
            return function_builder->build(arguments)->execute(arguments, result_type, input_rows_count);
        }

        ColumnPtr result_column;
        bool valid = castType(arguments[0].type.get(), [&](const auto & type)
        {
            using DataType = std::decay_t<decltype(type)>;

            if constexpr (std::is_same_v<DataTypeFixedString, DataType>)
            {
                if constexpr (allow_string_or_fixed_string)
                {
                    if (const auto * col = checkAndGetColumn<ColumnFixedString>(arguments[0].column.get()))
                    {
                        if constexpr (is_bit_count)
                        {
                            auto size = col->size();

                            auto col_res = ColumnUInt16::create(size);
                            auto & vec_res = col_res->getData();
                            vec_res.resize(col->size());

                            const auto & chars = col->getChars();
                            auto n = col->getN();
                            for (size_t i = 0; i < size; ++i)
                            {
                                vec_res[i] = StringUnaryOperationReduceImpl<Op<UInt8>>::vector(
                                    chars.data() + n * i, chars.data() + n * (i + 1));
                            }
                            result_column = std::move(col_res);
                            return true;
                        }
                        else
                        {
                            auto col_res = ColumnFixedString::create(col->getN());
                            auto & vec_res = col_res->getChars();
                            vec_res.resize(col->size() * col->getN());
                            FixedStringUnaryOperationImpl<Op<UInt8>>::vector(col->getChars(), vec_res);
                            result_column = std::move(col_res);
                            return true;
                        }
                    }
                }
            }
            else if constexpr (std::is_same_v<DataTypeString, DataType>)
            {
                if constexpr (allow_string_or_fixed_string)
                {
                    if (const auto * col = checkAndGetColumn<ColumnString>(arguments[0].column.get()))
                    {
                        if constexpr (is_bit_count)
                        {
                            auto size = col->size();

                            auto col_res = ColumnUInt64::create(size);
                            auto & vec_res = col_res->getData();

                            const auto & chars = col->getChars();
                            const auto & offsets = col->getOffsets();
                            for (size_t i = 0; i < size; ++i)
                            {
                                vec_res[i] = StringUnaryOperationReduceImpl<Op<UInt8>>::vector(
                                    chars.data() + offsets[i - 1], chars.data() + offsets[i] - 1);
                            }
                            result_column = std::move(col_res);
                            return true;
                        }
                        else
                        {
                            auto col_res = ColumnString::create();
                            auto & vec_res = col_res->getChars();
                            auto & offset_res = col_res->getOffsets();

                            const auto & vec_col = col->getChars();
                            const auto & offset_col = col->getOffsets();

                            vec_res.resize(vec_col.size());
                            offset_res.resize(offset_col.size());
                            memcpy(offset_res.data(), offset_col.data(), offset_res.size() * sizeof(UInt64));

                            FixedStringUnaryOperationImpl<Op<UInt8>>::vector(vec_col, vec_res);
                            result_column = std::move(col_res);
                            return true;
                        }
                    }
                }
            }
            else if constexpr (IsDataTypeDecimal<DataType>)
            {
                using T0 = typename DataType::FieldType;
                if constexpr (allow_decimal)
                {
                    if (auto col = checkAndGetColumn<ColumnDecimal<T0>>(arguments[0].column.get()))
                    {
                        if constexpr (is_sign_function)
                        {
                            auto col_res = ColumnVector<typename Op<T0>::ResultType>::create();
                            auto & vec_res = col_res->getData();
                            vec_res.resize(col->getData().size());
                            UnaryOperationImpl<T0, Op<T0>>::vector(col->getData(), vec_res);
                            result_column = std::move(col_res);
                            return true;
                        }
                        else
                        {
                            auto col_res = ColumnDecimal<typename Op<T0>::ResultType>::create(0, type.getScale());
                            auto & vec_res = col_res->getData();
                            vec_res.resize(col->getData().size());
                            UnaryOperationImpl<T0, Op<T0>>::vector(col->getData(), vec_res);
                            result_column = std::move(col_res);
                            return true;
                        }
                    }
                }
            }
            else
            {
                using T0 = typename DataType::FieldType;
                if (auto col = checkAndGetColumn<ColumnVector<T0>>(arguments[0].column.get()))
                {
                    auto col_res = ColumnVector<typename Op<T0>::ResultType>::create();
                    auto & vec_res = col_res->getData();
                    vec_res.resize(col->getData().size());
                    UnaryOperationImpl<T0, Op<T0>>::vector(col->getData(), vec_res);
                    result_column = std::move(col_res);
                    return true;
                }
            }

            return false;
        });
        if (!valid)
            throw Exception(ErrorCodes::LOGICAL_ERROR, "{}'s argument does not match the expected data type", getName());

        return result_column;
    }

#if USE_EMBEDDED_COMPILER
    bool isCompilableImpl(const DataTypes & arguments, const DataTypePtr & result_type) const override
    {
        if (1 != arguments.size())
            return false;

        if (!canBeNativeType(*arguments[0]) || !canBeNativeType(*result_type))
            return false;

        return castType(arguments[0].get(), [&](const auto & type)
        {
            using DataType = std::decay_t<decltype(type)>;
            if constexpr (std::is_same_v<DataTypeFixedString, DataType> || std::is_same_v<DataTypeString, DataType>)
            {
                return false;
            }
            else
            {
                using T0 = typename DataType::FieldType;
                using T1 = typename Op<T0>::ResultType;
                if constexpr (!std::is_same_v<T1, InvalidType> && !IsDataTypeDecimal<DataType> && Op<T0>::compilable)
                    return true;
            }

            return false;
        });
    }

    llvm::Value * compileImpl(llvm::IRBuilderBase & builder, const ValuesWithType & arguments, const DataTypePtr & result_type) const override
    {
        assert(1 == arguments.size());

        llvm::Value * result = nullptr;
        castType(arguments[0].type.get(), [&](const auto & type)
        {
            using DataType = std::decay_t<decltype(type)>;
            if constexpr (std::is_same_v<DataTypeFixedString, DataType> || std::is_same_v<DataTypeString, DataType>)
            {
                return false;
            }
            else
            {
                using T0 = typename DataType::FieldType;
                using T1 = typename Op<T0>::ResultType;
                if constexpr (!std::is_same_v<T1, InvalidType> && !IsDataTypeDecimal<DataType> && Op<T0>::compilable)
                {
                    auto & b = static_cast<llvm::IRBuilder<> &>(builder);
                    auto * v = nativeCast(b, arguments[0], result_type);
                    result = Op<T0>::compile(b, v, is_signed_v<T1>);

                    return true;
                }
            }

            return false;
        });

        return result;
    }
#endif

    bool hasInformationAboutMonotonicity() const override
    {
        return FunctionUnaryArithmeticMonotonicity<Name>::has();
    }

    Monotonicity getMonotonicityForRange(const IDataType &, const Field & left, const Field & right) const override
    {
        return FunctionUnaryArithmeticMonotonicity<Name>::get(left, right);
    }
};


struct PositiveMonotonicity
{
    static bool has() { return true; }
    static IFunction::Monotonicity get(const Field &, const Field &)
    {
        return { .is_monotonic = true };
    }
};

}