be/src/exprs/aggregate/aggregate_function_foreach.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.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/AggregateFunctions/Combinators/AggregateFunctionForEach.h
// and modified by Doris

#pragma once

#include <utility>

#include "common/status.h"
#include "core/assert_cast.h"
#include "core/column/column_nullable.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "exec/common/arithmetic_overflow.h"
#include "exprs/aggregate/aggregate_function.h"
#include "exprs/function/array/function_array_utils.h"

namespace doris {

struct AggregateFunctionForEachData {
    size_t dynamic_array_size = 0;
    char* array_of_aggregate_datas = nullptr;
};

/** Adaptor for aggregate functions.
  * Adding -ForEach suffix to aggregate function
  *  will convert that aggregate function to a function, accepting arrays,
  *  and applies aggregation for each corresponding elements of arrays independently,
  *  returning arrays of aggregated values on corresponding positions.
  *
  * Example: sumForEach of:
  *  [1, 2],
  *  [3, 4, 5],
  *  [6, 7]
  * will return:
  *  [10, 13, 5]
  *
  * TODO Allow variable number of arguments.
  */
class AggregateFunctionForEach : public AggregateFunctionNonFinalBase,
                                 public IAggregateFunctionDataHelper<AggregateFunctionForEachData,
                                                                     AggregateFunctionForEach>,
                                 VarargsExpression,
                                 NullableAggregateFunction {
protected:
    using Base =
            IAggregateFunctionDataHelper<AggregateFunctionForEachData, AggregateFunctionForEach>;

    AggregateFunctionPtr nested_function;
    const size_t nested_size_of_data;
    const size_t num_arguments;

    AggregateFunctionForEachData& ensure_aggregate_data(AggregateDataPtr __restrict place,
                                                        size_t new_size, Arena& arena) const {
        AggregateFunctionForEachData& state = data(place);

        /// Ensure we have aggregate states for new_size elements, allocate
        /// from arena if needed. When reallocating, we can't copy the
        /// states to new buffer with memcpy, because they may contain pointers
        /// to themselves. In particular, this happens when a state contains
        /// a PODArrayWithStackMemory, which stores small number of elements
        /// inline. This is why we create new empty states in the new buffer,
        /// and merge the old states to them.
        size_t old_size = state.dynamic_array_size;
        if (old_size < new_size) {
            static constexpr size_t MAX_ARRAY_SIZE = 100 * 1000000000ULL;
            if (new_size > MAX_ARRAY_SIZE) {
                throw Exception(ErrorCode::INTERNAL_ERROR,
                                "Suspiciously large array size ({}) in -ForEach aggregate function",
                                new_size);
            }

            size_t allocation_size = 0;
            if (common::mul_overflow(new_size, nested_size_of_data, allocation_size)) {
                throw Exception(ErrorCode::INTERNAL_ERROR,
                                "Allocation size ({} * {}) overflows in -ForEach aggregate "
                                "function, but it should've been prevented by previous checks",
                                new_size, nested_size_of_data);
            }

            char* old_state = state.array_of_aggregate_datas;

            char* new_state =
                    arena.aligned_alloc(allocation_size, nested_function->align_of_data());

            size_t num_created = 0;
            try {
                for (; num_created < new_size; ++num_created) {
                    nested_function->create(&new_state[num_created * nested_size_of_data]);
                }

                for (size_t i = 0; i < old_size; ++i) {
                    nested_function->merge(&new_state[i * nested_size_of_data],
                                           &old_state[i * nested_size_of_data], arena);
                }
            } catch (...) {
                for (size_t i = 0; i < num_created; ++i) {
                    nested_function->destroy(&new_state[i * nested_size_of_data]);
                }

                throw;
            }

            for (size_t i = 0; i < old_size; ++i) {
                nested_function->destroy(&old_state[i * nested_size_of_data]);
            }

            state.array_of_aggregate_datas = new_state;
            state.dynamic_array_size = new_size;
        }

        return state;
    }

private:
    struct PreparedBatchColumns {
        Columns owned_columns;
        std::vector<const IColumn*> nested_columns;
        const ColumnArray::Offsets64* offsets = nullptr;
    };

    void add_aligned(AggregateDataPtr __restrict place, const IColumn** nested,
                     const ColumnArray::Offsets64& offsets, size_t row_num, Arena& arena) const {
        const size_t begin = offsets[row_num - 1];
        const size_t end = offsets[row_num];
        AggregateFunctionForEachData& state = ensure_aggregate_data(place, end - begin, arena);

        char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = begin; i < end; ++i) {
            nested_function->add(nested_state, nested, i, arena);
            nested_state += nested_size_of_data;
        }
    }

    // Validate all selected rows first. If their physical begins differ, compact every argument
    // once for the batch and give them shared offsets;
    // E.g. data [9,8,7,6]/[1,2,3,4,5], offsets [1,4]/[2,5], skip row 0
    // -> data [8,7,6]/[3,4,5], shared offsets [0,3].
    template <typename IsSelected>
    Columns normalize_batch_columns(size_t batch_size, const IColumn** columns,
                                    IsSelected&& is_selected) const {
        bool offsets_aligned = true;
        // First pass: reject size mismatches before allocating, and detect whether the
        // original nested columns can use one shared element index.
        for (size_t row = 0; row < batch_size; ++row) {
            if (!is_selected(row)) {
                continue;
            }

            const auto& first_array =
                    assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[0]);
            const auto& first_offsets = first_array.get_offsets();
            const size_t first_begin = first_offsets[row - 1];
            const size_t row_size = first_offsets[row] - first_begin;
            for (size_t i = 1; i < num_arguments; ++i) {
                const auto& array =
                        assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[i]);
                const auto& offsets = array.get_offsets();
                const size_t begin = offsets[row - 1];
                if (offsets[row] - begin != row_size) {
                    throw Exception(ErrorCode::INTERNAL_ERROR,
                                    "Arrays passed to {} aggregate function have different sizes",
                                    get_name());
                }
                offsets_aligned &= row_size == 0 || begin == first_begin;
            }
        }

        // Fast path: all visible rows already address the same physical nested positions.
        if (offsets_aligned) {
            return {};
        }

        // Slow path: build the shared offsets once, copy only selected row slices, and preserve
        // the outer row count by emitting repeated offsets for unselected rows.
        auto offsets_column = ColumnArray::ColumnOffsets::create();
        auto& normalized_offsets = offsets_column->get_data();
        normalized_offsets.reserve(batch_size);
        size_t normalized_offset = 0;
        const auto& first_offsets =
                assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[0])
                        .get_offsets();
        for (size_t row = 0; row < batch_size; ++row) {
            if (is_selected(row)) {
                normalized_offset += first_offsets[row] - first_offsets[row - 1];
            }
            normalized_offsets.push_back(normalized_offset);
        }
        ColumnPtr shared_offsets = std::move(offsets_column);

        Columns normalized_columns(num_arguments);
        for (size_t i = 0; i < num_arguments; ++i) {
            const auto& array =
                    assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[i]);
            auto nested = array.get_data().clone_empty();
            for (size_t row = 0; row < batch_size; ++row) {
                if (is_selected(row)) {
                    const auto& offsets = array.get_offsets();
                    const size_t begin = offsets[row - 1];
                    const size_t row_size = offsets[row] - begin;
                    nested->insert_range_from(array.get_data(), begin, row_size);
                }
            }
            ColumnPtr nested_column = std::move(nested);
            normalized_columns[i] = ColumnArray::create(nested_column, shared_offsets);
        }
        return normalized_columns;
    }

    template <typename IsSelected>
    PreparedBatchColumns prepare_batch_columns(size_t batch_size, const IColumn** columns,
                                               IsSelected&& is_selected) const {
        PreparedBatchColumns prepared;
        prepared.owned_columns =
                normalize_batch_columns(batch_size, columns, std::forward<IsSelected>(is_selected));
        prepared.nested_columns.resize(num_arguments);
        for (size_t i = 0; i < num_arguments; ++i) {
            const IColumn* column =
                    prepared.owned_columns.empty() ? columns[i] : prepared.owned_columns[i].get();
            const auto& array =
                    assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*column);
            prepared.nested_columns[i] = &array.get_data();
            if (i == 0) {
                prepared.offsets = &array.get_offsets();
            }
        }
        return prepared;
    }

public:
    constexpr static auto AGG_FOREACH_SUFFIX = "_foreach";
    AggregateFunctionForEach(AggregateFunctionPtr nested_function_, const DataTypes& arguments)
            : Base(arguments),
              nested_function {std::move(nested_function_)},
              nested_size_of_data(nested_function->size_of_data()),
              num_arguments(arguments.size()) {
        if (arguments.empty()) {
            throw Exception(ErrorCode::INTERNAL_ERROR,
                            "Aggregate function {} require at least one argument", get_name());
        }
    }
    void set_version(const int version_) override {
        Base::set_version(version_);
        nested_function->set_version(version_);
    }

    String get_name() const override { return nested_function->get_name() + AGG_FOREACH_SUFFIX; }

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

    // Use foreach-specific streaming handling to align array inputs once while keeping aggregate
    // states row-local; prepare returns empty when the original columns are already aligned.
    bool requires_batch_add_for_streaming() const { return true; }

    template <typename IsSelected>
    Columns prepare_batch_columns_for_streaming(size_t batch_size, const IColumn** columns,
                                                IsSelected&& is_selected) const {
        return normalize_batch_columns(batch_size, columns, std::forward<IsSelected>(is_selected));
    }

    void destroy(AggregateDataPtr __restrict place) const noexcept override {
        AggregateFunctionForEachData& state = data(place);

        char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = 0; i < state.dynamic_array_size; ++i) {
            nested_function->destroy(nested_state);
            nested_state += nested_size_of_data;
        }
    }

    bool is_trivial() const override {
        return std::is_trivial_v<Data> && nested_function->is_trivial();
    }

    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
               Arena& arena) const override {
        const AggregateFunctionForEachData& rhs_state = data(rhs);
        AggregateFunctionForEachData& state =
                ensure_aggregate_data(place, rhs_state.dynamic_array_size, arena);

        const char* rhs_nested_state = rhs_state.array_of_aggregate_datas;
        char* nested_state = state.array_of_aggregate_datas;

        for (size_t i = 0; i < state.dynamic_array_size && i < rhs_state.dynamic_array_size; ++i) {
            nested_function->merge(nested_state, rhs_nested_state, arena);

            rhs_nested_state += nested_size_of_data;
            nested_state += nested_size_of_data;
        }
    }

    void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {
        const AggregateFunctionForEachData& state = data(place);
        buf.write_binary(state.dynamic_array_size);
        const char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = 0; i < state.dynamic_array_size; ++i) {
            nested_function->serialize(nested_state, buf);
            nested_state += nested_size_of_data;
        }
    }

    void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,
                     Arena& arena) const override {
        AggregateFunctionForEachData& state = data(place);

        size_t new_size = 0;
        buf.read_binary(new_size);

        ensure_aggregate_data(place, new_size, arena);

        char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = 0; i < new_size; ++i) {
            nested_function->deserialize(nested_state, buf, arena);
            nested_state += nested_size_of_data;
        }
    }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {
        const AggregateFunctionForEachData& state = data(place);

        auto& arr_to = assert_cast<ColumnArray&, TypeCheckOnRelease::DISABLE>(to);
        auto& offsets_to = arr_to.get_offsets();
        IColumn* elems_to = &arr_to.get_data();
        ColumnNullable* nullable_elems_to = nullptr;
        if (!nested_function->get_return_type()->is_nullable()) {
            nullable_elems_to = assert_cast<ColumnNullable*, TypeCheckOnRelease::DISABLE>(elems_to);
            elems_to = nullable_elems_to->get_nested_column_ptr().get();
        }

        char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = 0; i < state.dynamic_array_size; ++i) {
            nested_function->insert_result_into(nested_state, *elems_to);
            if (nullable_elems_to != nullptr) {
                nullable_elems_to->get_null_map_data().push_back(0);
            }
            nested_state += nested_size_of_data;
        }

        offsets_to.push_back(offsets_to.back() + state.dynamic_array_size);
    }

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

        const IColumn* elems_to = &arr_to->get_data();
        if (!nested_function->get_return_type()->is_nullable()) {
            const auto* nullable_elems_to = check_and_get_column<ColumnNullable>(*elems_to);
            if (UNLIKELY(nullable_elems_to == nullptr)) {
                throw doris::Exception(Status::InternalError(
                        "Aggregate function {} result type check failed: Array nested column "
                        "type {} is not ColumnNullable",
                        get_name(), elems_to->get_name()));
            }
            elems_to = &nullable_elems_to->get_nested_column();
        }
        nested_function->check_result_column_type(*elems_to);
    }

    void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
             Arena& arena) const override {
        std::vector<const IColumn*> nested(num_arguments);

        for (size_t i = 0; i < num_arguments; ++i) {
            nested[i] = &assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[i])
                                 .get_data();
        }

        const auto& first_array_column =
                assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[0]);
        const auto& offsets = first_array_column.get_offsets();

        size_t begin = offsets[row_num - 1];
        size_t end = offsets[row_num];
        const size_t row_size = end - begin;
        bool offsets_aligned = true;

        /// Sanity check. NOTE We can implement specialization for a case with single argument, if the check will hurt performance.
        for (size_t i = 1; i < num_arguments; ++i) {
            const auto& ith_column =
                    assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[i]);
            const auto& ith_offsets = ith_column.get_offsets();
            const size_t ith_begin = ith_offsets[row_num - 1];
            const size_t ith_end = ith_offsets[row_num];

            if (ith_end - ith_begin != row_size) {
                throw Exception(ErrorCode::INTERNAL_ERROR,
                                "Arrays passed to {} aggregate function have different sizes",
                                get_name());
            }
            offsets_aligned &= ith_begin == begin;
        }

        std::vector<ColumnPtr> compacted_nested;
        if (!offsets_aligned && row_size != 0) {
            // The nested aggregate accepts one shared index, so align only mismatched row slices.
            compacted_nested.reserve(num_arguments);
            for (size_t i = 0; i < num_arguments; ++i) {
                const auto& ith_column =
                        assert_cast<const ColumnArray&, TypeCheckOnRelease::DISABLE>(*columns[i]);
                const size_t ith_begin = ith_column.get_offsets()[row_num - 1];
                compacted_nested.emplace_back(nested[i]->cut(ith_begin, row_size));
                nested[i] = compacted_nested.back().get();
            }
            begin = 0;
            end = row_size;
        }

        AggregateFunctionForEachData& state = ensure_aggregate_data(place, row_size, arena);

        char* nested_state = state.array_of_aggregate_datas;
        for (size_t i = begin; i < end; ++i) {
            nested_function->add(nested_state, nested.data(), i, arena);
            nested_state += nested_size_of_data;
        }
    }

    void add_batch(size_t batch_size, AggregateDataPtr* places, size_t place_offset,
                   const IColumn** columns, Arena& arena, bool /*agg_many*/) const override {
        auto prepared = prepare_batch_columns(batch_size, columns, [](size_t) { return true; });
        for (size_t row = 0; row < batch_size; ++row) {
            add_aligned(places[row] + place_offset, prepared.nested_columns.data(),
                        *prepared.offsets, row, arena);
        }
    }

    void add_batch_selected(size_t batch_size, AggregateDataPtr* places, size_t place_offset,
                            const IColumn** columns, Arena& arena) const override {
        auto prepared = prepare_batch_columns(batch_size, columns,
                                              [&](size_t row) { return places[row] != nullptr; });
        for (size_t row = 0; row < batch_size; ++row) {
            if (places[row] != nullptr) {
                add_aligned(places[row] + place_offset, prepared.nested_columns.data(),
                            *prepared.offsets, row, arena);
            }
        }
    }

    void add_batch_single_place(size_t batch_size, AggregateDataPtr place, const IColumn** columns,
                                Arena& arena) const override {
        auto prepared = prepare_batch_columns(batch_size, columns, [](size_t) { return true; });
        for (size_t row = 0; row < batch_size; ++row) {
            add_aligned(place, prepared.nested_columns.data(), *prepared.offsets, row, arena);
        }
    }

    void check_input_columns_type(const IColumn** columns) const override {
        std::vector<const IColumn*> nested(num_arguments);
        for (size_t i = 0; i < num_arguments; ++i) {
            const auto* array_column = check_and_get_column<ColumnArray>(*columns[i]);
            if (UNLIKELY(array_column == nullptr)) {
                throw doris::Exception(Status::InternalError(
                        "Aggregate function {} argument {} type check failed: Column type {} is "
                        "not ColumnArray",
                        get_name(), i, columns[i]->get_name()));
            }
            nested[i] = &array_column->get_data();
        }
        nested_function->check_input_columns_type(nested.data());
    }
};
} // namespace doris

Coverage Report

Created: 2026-10-01 17:01