be/src/exprs/aggregate/aggregate_function_topn.h

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.

#pragma once

#include <rapidjson/encodings.h>
#include <rapidjson/stringbuffer.h>
#include <rapidjson/writer.h>
#include <stddef.h>
#include <stdint.h>

#include <algorithm>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "common/cast_set.h"
#include "common/exception.h"
#include "core/assert_cast.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
#include "core/column/column_decimal.h"
#include "core/column/column_nullable.h"
#include "core/column/column_string.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_string.h"
#include "core/string_buffer.hpp"
#include "core/string_ref.h"
#include "core/types.h"
#include "exec/common/hash_table/phmap_fwd_decl.h"
#include "exprs/aggregate/aggregate_function.h"
#include "exprs/aggregate/aggregate_function_simple_factory.h"

namespace doris {} // namespace doris

namespace doris {

// space-saving algorithm
template <PrimitiveType T>
struct AggregateFunctionTopNData {
    using ColVecType = typename PrimitiveTypeTraits<T>::ColumnType;
    using DataType = typename PrimitiveTypeTraits<T>::CppType;
    void set_paramenters(int input_top_num, int space_expand_rate = 50) {
        // FE validates a literal count, and a constant count only BE can evaluate is validated here.
        if (UNLIKELY(input_top_num <= 0)) {
            throw Exception(ErrorCode::INVALID_ARGUMENT,
                            "the count of topn, topn_array and topn_weighted must be a constant "
                            "positive integer, but got {}",
                            input_top_num);
        }
        top_num = input_top_num;
        // Non-positive expansion rates retain all candidates during serialization and merging.
        capacity = space_expand_rate <= 0 ? UINT64_MAX : (uint64_t)top_num * space_expand_rate;
    }

    void add(const StringRef& value, const UInt64& increment = 1) {
        std::string data = value.to_string();
        auto it = counter_map.find(data);
        if (it != counter_map.end()) {
            it->second = it->second + increment;
        } else {
            counter_map.insert({data, increment});
        }
    }

    void add(const DataType& value, const UInt64& increment = 1) {
        auto it = counter_map.find(value);
        if (it != counter_map.end()) {
            it->second = it->second + increment;
        } else {
            counter_map.insert({value, increment});
        }
    }

    void merge(const AggregateFunctionTopNData& rhs) {
        if (!rhs.top_num || rhs.counter_map.empty()) {
            return;
        }

        if (counter_map.empty()) {
            *this = rhs;
            return;
        }
        if (UNLIKELY(top_num != rhs.top_num || capacity != rhs.capacity)) {
            throw Exception(ErrorCode::INVALID_ARGUMENT,
                            "topn aggregate states have incompatible parameters: "
                            "({}, {}) vs ({}, {}) (N, capacity)",
                            top_num, capacity, rhs.top_num, rhs.capacity);
        }

        bool lhs_full = (counter_map.size() >= capacity);
        bool rhs_full = (rhs.counter_map.size() >= capacity);

        uint64_t lhs_min = 0;
        uint64_t rhs_min = 0;

        if (lhs_full) {
            lhs_min = UINT64_MAX;
            for (auto it : counter_map) {
                lhs_min = std::min(lhs_min, it.second);
            }
        }

        if (rhs_full) {
            rhs_min = UINT64_MAX;
            for (auto it : rhs.counter_map) {
                rhs_min = std::min(rhs_min, it.second);
            }

            for (auto& it : counter_map) {
                it.second += rhs_min;
            }
        }

        for (auto rhs_it : rhs.counter_map) {
            auto lhs_it = counter_map.find(rhs_it.first);
            if (lhs_it != counter_map.end()) {
                lhs_it->second += rhs_it.second - rhs_min;
            } else {
                counter_map.insert({rhs_it.first, rhs_it.second + lhs_min});
            }
        }
    }

    std::vector<std::pair<uint64_t, typename PrimitiveTypeTraits<T>::CppType>> get_remain_vector()
            const {
        std::vector<std::pair<uint64_t, typename PrimitiveTypeTraits<T>::CppType>> counter_vector;
        for (auto it : counter_map) {
            counter_vector.emplace_back(it.second, it.first);
        }
        std::sort(counter_vector.begin(), counter_vector.end(),
                  std::greater<std::pair<uint64_t, typename PrimitiveTypeTraits<T>::CppType>>());
        return counter_vector;
    }

    void write(BufferWritable& buf) const {
        buf.write_binary(top_num);
        buf.write_binary(capacity);

        uint64_t element_number = std::min(capacity, (uint64_t)counter_map.size());
        buf.write_binary(element_number);

        auto counter_vector = get_remain_vector();

        for (auto i = 0; i < element_number; i++) {
            auto element = counter_vector[i];
            buf.write_binary(element.second);
            buf.write_binary(element.first);
        }
    }

    void read(BufferReadable& buf) {
        buf.read_binary(top_num);
        buf.read_binary(capacity);

        uint64_t element_number = 0;
        buf.read_binary(element_number);

        counter_map.clear();
        std::pair<DataType, uint64_t> element;
        for (auto i = 0; i < element_number; i++) {
            buf.read_binary(element.first);
            buf.read_binary(element.second);
            counter_map.insert(element);
        }
    }

    std::string get() const {
        auto counter_vector = get_remain_vector();

        rapidjson::StringBuffer buffer;
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        writer.StartObject();
        for (int i = 0; i < std::min((int)counter_vector.size(), top_num); i++) {
            const auto& element = counter_vector[i];
            writer.Key(element.second.data(), cast_set<rapidjson::SizeType>(element.second.size()));
            writer.Uint64(element.first);
        }
        writer.EndObject();

        return buffer.GetString();
    }

    void insert_result_into(ColVecType& to) const {
        auto counter_vector = get_remain_vector();
        for (int i = 0; i < std::min((int)counter_vector.size(), top_num); i++) {
            const auto& element = counter_vector[i];
            if constexpr (is_string_type(T)) {
                to.insert_data(element.second.c_str(), element.second.length());
            } else {
                to.get_data().push_back(element.second);
            }
        }
    }

    void reset() {
        counter_map.clear();
        top_num = 0;
        capacity = 0;
    }

    int top_num = 0;
    uint64_t capacity = 0;
    flat_hash_map<DataType, uint64_t> counter_map;
};

struct AggregateFunctionTopNImplInt {
    using Data = AggregateFunctionTopNData<TYPE_STRING>;
    static void add(Data& __restrict place, const IColumn** columns, size_t row_num) {
        place.set_paramenters(
                assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[1])
                        ->get_element(row_num));
        place.add(assert_cast<const ColumnString&, TypeCheckOnRelease::DISABLE>(*columns[0])
                          .get_data_at(row_num));
    }
};

struct AggregateFunctionTopNImplIntInt {
    using Data = AggregateFunctionTopNData<TYPE_STRING>;
    static void add(Data& __restrict place, const IColumn** columns, size_t row_num) {
        place.set_paramenters(
                assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[1])
                        ->get_element(row_num),
                assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[2])
                        ->get_element(row_num));
        place.add(assert_cast<const ColumnString&, TypeCheckOnRelease::DISABLE>(*columns[0])
                          .get_data_at(row_num));
    }
};

//for topn_array agg
template <PrimitiveType T, bool has_default_param>
struct AggregateFunctionTopNImplArray {
    using Data = AggregateFunctionTopNData<T>;
    using ColVecType = typename PrimitiveTypeTraits<T>::ColumnType;
    static constexpr bool has_default_parameter = has_default_param;
    static constexpr bool is_weighted = false;
    static String get_name() { return "topn_array"; }
    static void add(AggregateFunctionTopNData<T>& __restrict place, const IColumn** columns,
                    size_t row_num) {
        if constexpr (has_default_param) {
            place.set_paramenters(
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[1])
                            ->get_element(row_num),
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[2])
                            ->get_element(row_num));

        } else {
            place.set_paramenters(
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[1])
                            ->get_element(row_num));
        }
        if constexpr (is_string_type(T)) {
            place.add(assert_cast<const ColumnString&, TypeCheckOnRelease::DISABLE>(*columns[0])
                              .get_data_at(row_num));
        } else {
            typename PrimitiveTypeTraits<T>::CppType val =
                    assert_cast<const ColVecType&, TypeCheckOnRelease::DISABLE>(*columns[0])
                            .get_data()[row_num];
            place.add(val);
        }
    }
};

//for topn_weighted agg
template <PrimitiveType T, bool has_default_param>
struct AggregateFunctionTopNImplWeight {
    using Data = AggregateFunctionTopNData<T>;
    using ColVecType = typename PrimitiveTypeTraits<T>::ColumnType;
    static constexpr bool has_default_parameter = has_default_param;
    static constexpr bool is_weighted = true;
    static String get_name() { return "topn_weighted"; }
    static void add(AggregateFunctionTopNData<T>& __restrict place, const IColumn** columns,
                    size_t row_num) {
        if constexpr (has_default_param) {
            place.set_paramenters(
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[2])
                            ->get_element(row_num),
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[3])
                            ->get_element(row_num));

        } else {
            place.set_paramenters(
                    assert_cast<const ColumnInt32*, TypeCheckOnRelease::DISABLE>(columns[2])
                            ->get_element(row_num));
        }
        if constexpr (is_string_type(T)) {
            auto weight = assert_cast<const ColumnInt64&, TypeCheckOnRelease::DISABLE>(*columns[1])
                                  .get_data()[row_num];
            place.add(assert_cast<const ColumnString&, TypeCheckOnRelease::DISABLE>(*columns[0])
                              .get_data_at(row_num),
                      weight);
        } else {
            typename PrimitiveTypeTraits<T>::CppType val =
                    assert_cast<const typename PrimitiveTypeTraits<T>::ColumnType&,
                                TypeCheckOnRelease::DISABLE>(*columns[0])
                            .get_data()[row_num];
            auto weight = assert_cast<const ColumnInt64&, TypeCheckOnRelease::DISABLE>(*columns[1])
                                  .get_data()[row_num];
            place.add(val, weight);
        }
    }
};

//base function
template <typename Impl, typename Derived>
class AggregateFunctionTopNBase
        : public IAggregateFunctionDataHelper<typename Impl::Data, Derived> {
public:
    AggregateFunctionTopNBase(const DataTypes& argument_types_)
            : IAggregateFunctionDataHelper<typename Impl::Data, Derived>(argument_types_) {}

    void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
             Arena&) const override {
        Impl::add(this->data(place), columns, row_num);
    }

    void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }

    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
               Arena&) const override {
        this->data(place).merge(this->data(rhs));
    }

    void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {
        this->data(place).write(buf);
    }

    void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,
                     Arena&) const override {
        this->data(place).read(buf);
    }
};

//topn function return string
template <typename Impl>
class AggregateFunctionTopN final
        : public AggregateFunctionTopNBase<Impl, AggregateFunctionTopN<Impl>>,
          MultiExpression,
          NullableAggregateFunction {
public:
    AggregateFunctionTopN(const DataTypes& argument_types_)
            : AggregateFunctionTopNBase<Impl, AggregateFunctionTopN<Impl>>(argument_types_) {}

    String get_name() const override { return "topn"; }

    DataTypePtr get_return_type() const override { return std::make_shared<DataTypeString>(); }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {
        std::string result = this->data(place).get();
        assert_cast<ColumnString&, TypeCheckOnRelease::DISABLE>(to).insert_data(result.c_str(),
                                                                                result.length());
    }

    void check_input_columns_type(const IColumn** columns) const override {
        this->template check_argument_column_type<ColumnString>(columns[0]);
        this->template check_argument_column_type<ColumnInt32>(columns[1]);
        if (this->argument_types.size() == 3) {
            this->template check_argument_column_type<ColumnInt32>(columns[2]);
        }
    }

    void check_result_column_type(const IColumn& column) const override {
        this->template check_result_column_type_as<ColumnString>(column);
    }
};

//topn function return array
template <typename Impl>
class AggregateFunctionTopNArray final
        : public AggregateFunctionTopNBase<Impl, AggregateFunctionTopNArray<Impl>>,
          MultiExpression,
          NullableAggregateFunction {
private:
    using ColVecType = typename Impl::Data::ColVecType;

public:
    AggregateFunctionTopNArray(const DataTypes& argument_types_)
            : AggregateFunctionTopNBase<Impl, AggregateFunctionTopNArray<Impl>>(argument_types_),
              _argument_type(argument_types_[0]) {}

    String get_name() const override { return Impl::get_name(); }

    DataTypePtr get_return_type() const override {
        return std::make_shared<DataTypeArray>(make_nullable(_argument_type));
    }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {
        auto& to_arr = assert_cast<ColumnArray&, TypeCheckOnRelease::DISABLE>(to);
        auto& col_null =
                assert_cast<ColumnNullable&, TypeCheckOnRelease::DISABLE>(to_arr.get_data());
        auto& typed_to =
                assert_cast<ColVecType&, TypeCheckOnRelease::DISABLE>(col_null.get_nested_column());
        this->data(place).insert_result_into(typed_to);
        col_null.get_null_map_data().resize_fill(typed_to.size(), 0);
        to_arr.get_offsets().push_back(typed_to.size());
    }

    void check_input_columns_type(const IColumn** columns) const override {
        this->template check_argument_column_type<ColVecType>(columns[0]);
        if constexpr (Impl::is_weighted) {
            this->template check_argument_column_type<ColumnInt64>(columns[1]);
            this->template check_argument_column_type<ColumnInt32>(columns[2]);
            if constexpr (Impl::has_default_parameter) {
                this->template check_argument_column_type<ColumnInt32>(columns[3]);
            }
        } else {
            this->template check_argument_column_type<ColumnInt32>(columns[1]);
            if constexpr (Impl::has_default_parameter) {
                this->template check_argument_column_type<ColumnInt32>(columns[2]);
            }
        }
    }

    void check_result_column_type(const IColumn& to) const override {
        const auto* to_arr = check_and_get_column<ColumnArray>(to);
        if (UNLIKELY(to_arr == nullptr)) {
            throw doris::Exception(Status::InternalError(
                    "Aggregate function {} result type check failed: Column type {} ({}) is not "
                    "ColumnArray",
                    get_name(), to.get_name(), typeid(to).name()));
        }

        const IColumn& data_column = to_arr->get_data();
        const auto* nullable_column = check_and_get_column<ColumnNullable>(data_column);
        if (UNLIKELY(nullable_column == nullptr)) {
            throw doris::Exception(Status::InternalError(
                    "Aggregate function {} result type check failed: Array nested column type {} "
                    "({}) is not ColumnNullable",
                    get_name(), data_column.get_name(), typeid(data_column).name()));
        }

        const IColumn& nested_column = nullable_column->get_nested_column();
        if (UNLIKELY(check_and_get_column<ColVecType>(nested_column) == nullptr)) {
            throw doris::Exception(Status::InternalError(
                    "Aggregate function {} result type check failed: Array nested data column "
                    "type {} ({}) does not match expected physical column type {}",
                    get_name(), nested_column.get_name(), typeid(nested_column).name(),
                    typeid(ColVecType).name()));
        }
    }

    DataTypePtr _argument_type;
};

} // namespace doris

Coverage Report

Created: 2026-10-13 23:46