be/src/exprs/function/function_bitmap_variadic.cpp

Source
// 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, COUNT_OP)                        \
    struct CLASS {                                                                                \
    public:                                                                                       \
        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) { \
            DCHECK_GE(col_size, 2);                                                               \
            if (col_size == 2) {                                                                  \
                const auto lhs = make_bitmap_column_accessor(argument_columns[0]);                \
                const auto rhs = make_bitmap_column_accessor(argument_columns[1]);                \
                for (size_t row = 0; row < input_rows_count; ++row) {                             \
                    const auto* lhs_value = lhs.is_null_at(row) ? nullptr : &lhs.get_value(row);  \
                    const auto* rhs_value = rhs.is_null_at(row) ? nullptr : &rhs.get_value(row);  \
                    if (!lhs_value || !rhs_value) {                                               \
                        if constexpr (std::is_same_v<CLASS, BitmapOrCount>) {                     \
                            const auto* value = lhs_value ? lhs_value : rhs_value;                \
                            res[row] = value ? value->cardinality() : 0;                          \
                        } else {                                                                  \
                            res[row] = 0;                                                         \
                        }                                                                         \
                    } else {                                                                      \
                        res[row] = lhs_value->COUNT_OP(*rhs_value);                               \
                    }                                                                             \
                }                                                                                 \
                return Status::OK();                                                              \
            }                                                                                     \
            TData vals(input_rows_count);                                                         \
            const auto first_column = make_bitmap_column_accessor(argument_columns[0]);           \
            for (size_t row = 0; row < input_rows_count; ++row) {                                 \
                if (!first_column.is_null_at(row)) {                                              \
                    vals[row] = first_column.get_value(row);                                      \
                }                                                                                 \
            }                                                                                     \
            for (size_t col = 1; col + 1 < col_size; ++col) {                                     \
                const auto column = make_bitmap_column_accessor(argument_columns[col]);           \
                for (size_t row = 0; row < input_rows_count; ++row) {                             \
                    if (!column.is_null_at(row)) {                                                \
                        vals[row] OP column.get_value(row);                                       \
                    }                                                                             \
                }                                                                                 \
            }                                                                                     \
            const auto last_column = make_bitmap_column_accessor(argument_columns[col_size - 1]); \
            for (size_t row = 0; row < input_rows_count; ++row) {                                 \
                const auto* lhs_value = &vals[row];                                               \
                const auto* rhs_value =                                                           \
                        last_column.is_null_at(row) ? nullptr : &last_column.get_value(row);      \
                if (!rhs_value) {                                                                 \
                    if constexpr (std::is_same_v<CLASS, BitmapOrCount>) {                         \
                        res[row] = lhs_value->cardinality();                                      \
                    } else {                                                                      \
                        res[row] = 0;                                                             \
                    }                                                                             \
                } else {                                                                          \
                    res[row] = lhs_value->COUNT_OP(*rhs_value);                                   \
                }                                                                                 \
            }                                                                                     \
            return Status::OK();                                                                  \
        }                                                                                         \
                                                                                                  \
    private:                                                                                      \
        struct BitmapColumnAccessor {                                                             \
            const std::vector<BitmapValue>* values = nullptr;                                     \
            const BitmapValue* const_value = nullptr;                                             \
            const ColumnUInt8::value_type* null_map_data = nullptr;                               \
            bool is_const = false;                                                                \
            bool is_const_null = false;                                                           \
                                                                                                  \
            bool is_null_at(size_t row) const {                                                   \
                return is_const ? is_const_null : (null_map_data && null_map_data[row]);          \
            }                                                                                     \
                                                                                                  \
            const BitmapValue& get_value(size_t row) const {                                      \
                return is_const ? *const_value : (*values)[row];                                  \
            }                                                                                     \
        };                                                                                        \
                                                                                                  \
        static BitmapColumnAccessor make_bitmap_column_accessor(const ColumnPtr& column) {        \
            BitmapColumnAccessor accessor;                                                        \
            const auto& [data_column_ptr, is_const] = unpack_if_const(column);                    \
            accessor.is_const = is_const;                                                         \
            const IColumn* data_column = data_column_ptr.get();                                   \
                                                                                                  \
            if (const auto* nullable = check_and_get_column<ColumnNullable>(*data_column)) {      \
                if (accessor.is_const) {                                                          \
                    accessor.is_const_null = nullable->is_null_at(0);                             \
                } else {                                                                          \
                    accessor.null_map_data = nullable->get_null_map_data().data();                \
                }                                                                                 \
                data_column = nullable->get_nested_column_ptr().get();                            \
            }                                                                                     \
                                                                                                  \
            const auto* bitmap_column = assert_cast<const ColumnBitmap*>(data_column);            \
            if (accessor.is_const) {                                                              \
                accessor.const_value = &bitmap_column->get_data()[0];                             \
            } else {                                                                              \
                accessor.values = &bitmap_column->get_data();                                     \
            }                                                                                     \
                                                                                                  \
            return accessor;                                                                      \
        }                                                                                         \
    }

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, |=, or_cardinality);
BITMAP_FUNCTION_COUNT_VARIADIC(BitmapAndCount, bitmap_and_count, &=, and_cardinality);
BITMAP_FUNCTION_COUNT_VARIADIC(BitmapXorCount, bitmap_xor_count, ^=, xor_cardinality);

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, bool NewVersion = false>
class FunctionBitMapVariadic : public IFunction {
    static_assert(!NewVersion || (NewVersion && (std::is_same_v<Impl, BitmapOrCount> ||
                                                 std::is_same_v<Impl, BitmapAndCount> ||
                                                 std::is_same_v<Impl, BitmapXorCount>)));

public:
    static constexpr auto name = []() constexpr {
        if constexpr (!NewVersion) {
            return Impl::name;
        } else if constexpr (std::is_same_v<Impl, BitmapOrCount>) {
            return "bitmap_or_count_v2";
        } else if constexpr (std::is_same_v<Impl, BitmapAndCount>) {
            return "bitmap_and_count_v2";
        } else if constexpr (std::is_same_v<Impl, BitmapXorCount>) {
            return "bitmap_xor_count_v2";
        } else {
            return Impl::name;
        }
    }();

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

    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 constexpr (NewVersion) {
            return std::make_shared<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 constexpr (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();

        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);

        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;
            if constexpr (!std::is_same_v<Impl, BitmapOrCount> &&
                          !std::is_same_v<Impl, BitmapAndCount> &&
                          !std::is_same_v<Impl, BitmapXorCount>) {
                argument_columns[i] = argument_columns[i]->convert_to_full_column_if_const();
            }
        }
        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>;
using FunctionBitmapOrCountV2 = FunctionBitMapVariadic<BitmapOrCount, true>;
using FunctionBitmapAndCountV2 = FunctionBitMapVariadic<BitmapAndCount, true>;
using FunctionBitmapXorCountV2 = FunctionBitMapVariadic<BitmapXorCount, true>;

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>();
    factory.register_function<FunctionBitmapOrCountV2>();
    factory.register_function<FunctionBitmapAndCountV2>();
    factory.register_function<FunctionBitmapXorCountV2>();
}
} // namespace doris

Coverage Report

Created: 2026-10-09 18:00