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

// and modified by Doris

#pragma once

#include <stddef.h>

#include <algorithm>

#include <boost/iterator/iterator_facade.hpp>

#include <memory>

#include <vector>

#include "core/assert_cast.h"

#include "core/column/column.h"

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

#include "core/data_type/data_type_number.h"

#include "core/types.h"

#include "exprs/aggregate/aggregate_function.h"

#include "util/simd/bits.h"

namespace doris {

class Arena;

class BufferReadable;

class BufferWritable;

struct AggregateFunctionCountData {

    UInt64 count = 0;

};

/// Simply count number of calls.

class AggregateFunctionCount final

        : public IAggregateFunctionDataHelper<AggregateFunctionCountData, AggregateFunctionCount>,

          VarargsExpression,

          NotNullableAggregateFunction {

public:

    AggregateFunctionCount(const DataTypes& argument_types_)

            : IAggregateFunctionDataHelper(argument_types_) {}

    bool is_simple_count() const override { return true; }

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

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

    bool is_trivial() const override { return true; }

    void add(AggregateDataPtr __restrict place, const IColumn**, ssize_t, Arena&) const override {

        ++data(place).count;

    }

    void add_batch_single_place(size_t batch_size, AggregateDataPtr place, const IColumn**,

                                Arena&) const override {

        data(place).count += batch_size;

    }

    void reset(AggregateDataPtr place) const override {

        AggregateFunctionCount::data(place).count = 0;

    }

    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,

               Arena&) const override {

        data(place).count += data(rhs).count;

    }

    void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {

        buf.write_var_uint(data(place).count);

    }

    void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,

                     Arena&) const override {

        buf.read_var_uint(data(place).count);

    }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {

        assert_cast<ColumnInt64&, TypeCheckOnRelease::DISABLE>(to).get_data().push_back(

                data(place).count);

    }

    void serialize_to_column(const std::vector<AggregateDataPtr>& places, size_t offset,

                             MutableColumnPtr& dst, const size_t num_rows) const override {

        auto& col = assert_cast<ColumnFixedLengthObject&>(*dst);

        DCHECK(col.item_size() == sizeof(Data))

                << "size is not equal: " << col.item_size() << " " << sizeof(Data);

        col.resize(num_rows);

        auto* data = col.get_data().data();

        for (size_t i = 0; i != num_rows; ++i) {

            *reinterpret_cast<Data*>(&data[sizeof(Data) * i]) =

                    *reinterpret_cast<Data*>(places[i] + offset);

        }

    }

    void streaming_agg_serialize_to_column(const IColumn** columns, MutableColumnPtr& dst,

                                           const size_t num_rows, Arena&) const override {

        auto& dst_col = assert_cast<ColumnFixedLengthObject&>(*dst);

        DCHECK(dst_col.item_size() == sizeof(Data))

                << "size is not equal: " << dst_col.item_size() << " " << sizeof(Data);

        dst_col.resize(num_rows);

        auto* data = dst_col.get_data().data();

        for (size_t i = 0; i != num_rows; ++i) {

            auto& state = *reinterpret_cast<Data*>(&data[sizeof(Data) * i]);

            state.count = 1;

        }

    }

    void deserialize_from_column_row(AggregateDataPtr __restrict place, const IColumn& column,

                                     size_t row, Arena&) const override {

        DCHECK_LT(row, column.size());

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(column);

        const auto* data = reinterpret_cast<const Data*>(col.get_data().data());

        AggregateFunctionCount::data(place) = data[row];

    }

    void deserialize_and_merge_from_column_range(AggregateDataPtr __restrict place,

                                                 const IColumn& column, size_t begin, size_t end,

                                                 Arena&) const override {

        DCHECK(end <= column.size() && begin <= end)

                << ", begin:" << begin << ", end:" << end << ", column.size():" << column.size();

        auto& col = assert_cast<const ColumnFixedLengthObject&>(column);

        auto* data = reinterpret_cast<const Data*>(col.get_data().data());

        for (size_t i = begin; i <= end; ++i) {

            doris::AggregateFunctionCount::data(place).count += data[i].count;

        }

    }

    bool needs_deserialize_and_merge_scratch() const override { return false; }

    void deserialize_and_merge_vec(const AggregateDataPtr* places, size_t offset,

                                   AggregateDataPtr rhs, const IColumn* column, Arena& arena,

                                   const size_t num_rows) const override {

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(*column);

        const auto* data = col.get_data().data();

        this->merge_vec(places, offset, AggregateDataPtr(data), arena, num_rows);

    }

    void deserialize_and_merge_vec_selected(const AggregateDataPtr* places, size_t offset,

                                            AggregateDataPtr rhs, const IColumn* column,

                                            Arena& arena, const size_t num_rows) const override {

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(*column);

        const auto* data = col.get_data().data();

        this->merge_vec_selected(places, offset, AggregateDataPtr(data), arena, num_rows);

    }

    void serialize_without_key_to_column(ConstAggregateDataPtr __restrict place,

                                         IColumn& to) const override {

        auto& col = assert_cast<ColumnFixedLengthObject&>(to);

        DCHECK(col.item_size() == sizeof(Data))

                << "size is not equal: " << col.item_size() << " " << sizeof(Data);

        size_t old_size = col.size();

        col.resize(old_size + 1);

        (reinterpret_cast<Data*>(col.get_data().data()) + old_size)->count =

                AggregateFunctionCount::data(place).count;

    }

    MutableColumnPtr create_serialize_column() const override {

        return ColumnFixedLengthObject::create(sizeof(Data));

    }

    DataTypePtr get_serialized_type() const override {

        return std::make_shared<DataTypeFixedLengthObject>();

    }

    void add_range_single_place(int64_t partition_start, int64_t partition_end, int64_t frame_start,

                                int64_t frame_end, AggregateDataPtr place, const IColumn** columns,

                                Arena& arena, UInt8* use_null_result,

                                UInt8* could_use_previous_result) const override {

        frame_start = std::max<int64_t>(frame_start, partition_start);

        frame_end = std::min<int64_t>(frame_end, partition_end);

        if (frame_start >= frame_end) {

            if (!*could_use_previous_result) {

                *use_null_result = true;

            }

        } else {

            AggregateFunctionCount::data(place).count += frame_end - frame_start;

            *use_null_result = false;

            *could_use_previous_result = true;

        }

    }

};

// Used for unary count(nullable_expr). SQL count(expr) counts non-NULL values.

class AggregateFunctionCountNotNullUnary final

        : public IAggregateFunctionDataHelper<AggregateFunctionCountData,

                                              AggregateFunctionCountNotNullUnary> {

public:

    AggregateFunctionCountNotNullUnary(const DataTypes& argument_types_)

            : IAggregateFunctionDataHelper(argument_types_) {}

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

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

    bool is_trivial() const override { return true; }

    void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,

             Arena&) const override {

        data(place).count +=

                !assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0])

                         .is_null_at(row_num);

    }

    void add_batch_single_place(size_t batch_size, AggregateDataPtr place, const IColumn** columns,

                                Arena&) const override {

        const auto& nullable_column =

                assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0]);

        const auto& null_map = nullable_column.get_null_map_data();

        DCHECK_LE(batch_size, null_map.size());

        if (!nullable_column.has_null(0, batch_size)) {

            data(place).count += batch_size;

            return;

        }

        data(place).count +=

                simd::count_zero_num(reinterpret_cast<const int8_t*>(null_map.data()), batch_size);

    }

    void reset(AggregateDataPtr place) const override { data(place).count = 0; }

    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,

               Arena&) const override {

        data(place).count += data(rhs).count;

    }

    void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {

        buf.write_var_uint(data(place).count);

    }

    void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,

                     Arena&) const override {

        buf.read_var_uint(data(place).count);

    }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {

        if (is_column_nullable(to)) {

            auto& null_column = assert_cast<ColumnNullable&, TypeCheckOnRelease::DISABLE>(to);

            null_column.get_null_map_data().push_back(0);

            assert_cast<ColumnInt64&, TypeCheckOnRelease::DISABLE>(null_column.get_nested_column())

                    .get_data()

                    .push_back(data(place).count);

        } else {

            assert_cast<ColumnInt64&, TypeCheckOnRelease::DISABLE>(to).get_data().push_back(

                    data(place).count);

        }

    }

    void check_result_column_type(const IColumn& to) const override {

        if (const auto* null_column = check_and_get_column<ColumnNullable>(to)) {

            IAggregateFunction::check_result_column_type(null_column->get_nested_column());

            return;

        }

        IAggregateFunction::check_result_column_type(to);

    }

    void serialize_to_column(const std::vector<AggregateDataPtr>& places, size_t offset,

                             MutableColumnPtr& dst, const size_t num_rows) const override {

        auto& col = assert_cast<ColumnFixedLengthObject&>(*dst);

        DCHECK(col.item_size() == sizeof(Data))

                << "size is not equal: " << col.item_size() << " " << sizeof(Data);

        col.resize(num_rows);

        auto* data = col.get_data().data();

        for (size_t i = 0; i != num_rows; ++i) {

            *reinterpret_cast<Data*>(&data[sizeof(Data) * i]) =

                    *reinterpret_cast<Data*>(places[i] + offset);

        }

    }

    void streaming_agg_serialize_to_column(const IColumn** columns, MutableColumnPtr& dst,

                                           const size_t num_rows, Arena&) const override {

        auto& col = assert_cast<ColumnFixedLengthObject&>(*dst);

        DCHECK(col.item_size() == sizeof(Data))

                << "size is not equal: " << col.item_size() << " " << sizeof(Data);

        col.resize(num_rows);

        auto& data = col.get_data();

        const ColumnNullable& input_col = assert_cast<const ColumnNullable&>(*columns[0]);

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

            auto& state = *reinterpret_cast<Data*>(&data[sizeof(Data) * i]);

            state.count = !input_col.is_null_at(i);

        }

    }

    void deserialize_from_column_row(AggregateDataPtr __restrict place, const IColumn& column,

                                     size_t row, Arena&) const override {

        DCHECK_LT(row, column.size());

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(column);

        const auto* data = reinterpret_cast<const Data*>(col.get_data().data());

        AggregateFunctionCountNotNullUnary::data(place) = data[row];

    }

    void deserialize_and_merge_from_column_range(AggregateDataPtr __restrict place,

                                                 const IColumn& column, size_t begin, size_t end,

                                                 Arena&) const override {

        DCHECK(end <= column.size() && begin <= end)

                << ", begin:" << begin << ", end:" << end << ", column.size():" << column.size();

        auto& col = assert_cast<const ColumnFixedLengthObject&>(column);

        auto* data = reinterpret_cast<const Data*>(col.get_data().data());

        for (size_t i = begin; i <= end; ++i) {

            doris::AggregateFunctionCountNotNullUnary::data(place).count += data[i].count;

        }

    }

    bool needs_deserialize_and_merge_scratch() const override { return false; }

    void deserialize_and_merge_vec(const AggregateDataPtr* places, size_t offset,

                                   AggregateDataPtr rhs, const IColumn* column, Arena& arena,

                                   const size_t num_rows) const override {

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(*column);

        const auto* data = col.get_data().data();

        this->merge_vec(places, offset, AggregateDataPtr(data), arena, num_rows);

    }

    void deserialize_and_merge_vec_selected(const AggregateDataPtr* places, size_t offset,

                                            AggregateDataPtr rhs, const IColumn* column,

                                            Arena& arena, const size_t num_rows) const override {

        const auto& col = assert_cast<const ColumnFixedLengthObject&>(*column);

        const auto* data = col.get_data().data();

        this->merge_vec_selected(places, offset, AggregateDataPtr(data), arena, num_rows);

    }

    void serialize_without_key_to_column(ConstAggregateDataPtr __restrict place,

                                         IColumn& to) const override {

        auto& col = assert_cast<ColumnFixedLengthObject&>(to);

        DCHECK(col.item_size() == sizeof(Data))

                << "size is not equal: " << col.item_size() << " " << sizeof(Data);

        size_t old_size = col.size();

        col.resize(old_size + 1);

        (reinterpret_cast<Data*>(col.get_data().data()) + old_size)->count =

                AggregateFunctionCountNotNullUnary::data(place).count;

    }

    MutableColumnPtr create_serialize_column() const override {

        return ColumnFixedLengthObject::create(sizeof(Data));

    }

    DataTypePtr get_serialized_type() const override {

        return std::make_shared<DataTypeFixedLengthObject>();

    }

    void add_range_single_place(int64_t partition_start, int64_t partition_end, int64_t frame_start,

                                int64_t frame_end, AggregateDataPtr place, const IColumn** columns,

                                Arena& arena, UInt8* use_null_result,

                                UInt8* could_use_previous_result) const override {

        frame_start = std::max<int64_t>(frame_start, partition_start);

        frame_end = std::min<int64_t>(frame_end, partition_end);

        if (frame_start >= frame_end) {

            if (!*could_use_previous_result) {

                *use_null_result = true;

            }

        } else {

            const auto& nullable_column =

                    assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0]);

            size_t count = 0;

            if (nullable_column.has_null()) {

                for (int64_t i = frame_start; i < frame_end; ++i) {

                    if (!nullable_column.is_null_at(i)) {

                        ++count;

                    }

                }

            } else {

                count = frame_end - frame_start;

            }

            *use_null_result = false;

            *could_use_previous_result = true;

            AggregateFunctionCountNotNullUnary::data(place).count += count;

        }

    }

};

} // namespace doris