[Feature](agg) add agg function entropy

be/src/vec/aggregate_functions/aggregate_function_entropy.cpp

@@ -0,0 +1,65 @@
+// 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/AggregateFunctionEntropy.cpp
+// and modified by Doris
+
+#include "vec/aggregate_functions/aggregate_function_entropy.h"
+
+#include "runtime/define_primitive_type.h"
+#include "vec/aggregate_functions/aggregate_function_simple_factory.h"
+#include "vec/aggregate_functions/helpers.h"
+
+namespace doris::vectorized {
+#include "common/compile_check_begin.h"
+
+AggregateFunctionPtr create_aggregate_function_entropy(const std::string& name,
+                                                       const DataTypes& argument_types,
+                                                       const DataTypePtr& result_type,
+                                                       const bool result_is_nullable,
+                                                       const AggregateFunctionAttr& attr) {
+    if (argument_types.size() == 1) {
+        auto res = creator_with_type_list<
+                TYPE_BOOLEAN, TYPE_TINYINT, TYPE_SMALLINT, TYPE_INT, TYPE_BIGINT, TYPE_LARGEINT,
+                TYPE_DECIMAL32, TYPE_DECIMAL64, TYPE_DECIMAL128I, TYPE_DECIMAL256, TYPE_DECIMALV2,
+                TYPE_FLOAT, TYPE_DOUBLE, TYPE_DATE, TYPE_DATETIME, TYPE_DATEV2, TYPE_DATETIMEV2,
+                TYPE_TIME, TYPE_TIMEV2, TYPE_TIMESTAMPTZ>::
+                create<AggregateFunctionEntropy, AggregateFunctionEntropySingleNumericData>(
+                        argument_types, result_is_nullable, attr);
+        if (res) {
+            return res;
+        }
+
+        auto type = argument_types[0]->get_primitive_type();
+        if (is_string_type(type) || is_varbinary(type) || type == TYPE_JSONB) {
+            res = creator_without_type::create<
+                    AggregateFunctionEntropy<AggregateFunctionEntropySingleStringData>>(
+                    argument_types, result_is_nullable, attr);
+            return res;
+        }
+    }
+
+    return creator_without_type::create<
+            AggregateFunctionEntropy<AggregateFunctionEntropyGenericData>>(
+            argument_types, result_is_nullable, attr);
+}
+
+void register_aggregate_function_entropy(AggregateFunctionSimpleFactory& factory) {
+    factory.register_function_both("entropy", create_aggregate_function_entropy);
+}
+
+} // namespace doris::vectorized

be/src/vec/aggregate_functions/aggregate_function_entropy.h

@@ -0,0 +1,191 @@
+// 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/AggregateFunctionEntropy.cpp
+// and modified by Doris
+
+#pragma once
+
+#include "vec/aggregate_functions/aggregate_function.h"
+#include "vec/columns/column.h"
+#include "vec/columns/column_decimal.h"
+#include "vec/common/assert_cast.h"
+#include "vec/common/hash_table/hash.h"
+#include "vec/common/hash_table/phmap_fwd_decl.h"
+#include "vec/common/string_ref.h"
+#include "vec/common/uint128.h"
+#include "vec/core/types.h"
+#include "vec/data_types/data_type.h"
+#include "vec/data_types/data_type_number.h"
+
+namespace doris::vectorized {
+#include "common/compile_check_begin.h"
+
+class Arena;
+class BufferReadable;
+class BufferWritable;
+template <PrimitiveType>
+class ColumnVector;
+
+/** Calculates Shannon Entropy, using HashMap and computing empirical distribution function.
+  * Entropy is measured in bits (base-2 logarithm is used).
+  */
+template <typename Value, typename Hash = HashCRC32<Value>>
+struct AggregateFunctionEntropyData {
+    using Container = flat_hash_map<Value, uint64_t, Hash>;
+    using Self = AggregateFunctionEntropyData<Value, Hash>;
+    Container frequency_map;
+
+    void clear() { frequency_map.clear(); }
+
+    void add(const Value& elem) { ++frequency_map[elem]; }
+
+    void merge(const Self& rhs) {
+        frequency_map.reserve(frequency_map.size() + rhs.frequency_map.size());
+        for (const auto& [elem, count] : rhs.frequency_map) {
+            frequency_map[elem] += count;
+        }
+    }
+
+    void write(BufferWritable& buf) const {
+        buf.write_var_uint(frequency_map.size());
+        for (const auto& [elem, count] : frequency_map) {
+            buf.write_binary(elem);
+            buf.write_binary(count);
+        }
+    }
+
+    void read(BufferReadable& buf) {
+        uint64_t new_size = 0;
+        buf.read_var_uint(new_size);
+        frequency_map.reserve(frequency_map.size() + new_size);
+
+        Value elem;
+        uint64_t count;
+        for (size_t i = 0; i < new_size; ++i) {
+            buf.read_binary(elem);
+            buf.read_binary(count);
+            frequency_map[elem] += count;
+        }
+    }
+
+    Float64 get_result() const {
+        Float64 entropy = 0;
+        uint64_t total_count = 0;
+        for (const auto& [_, count] : frequency_map) {
+            total_count += count;
+        }
+        for (const auto& [_, count] : frequency_map) {
+            Float64 p = static_cast<Float64>(count) / static_cast<Float64>(total_count);
+            entropy -= p * std::log2(p);
+        }
+        return entropy;
+    }
+
+    static String get_name() { return "entropy"; }
+};
+
+template <PrimitiveType T>
+struct AggregateFunctionEntropySingleNumericData
+        : public AggregateFunctionEntropyData<typename PrimitiveTypeTraits<T>::CppType> {
+    using Base = AggregateFunctionEntropyData<typename PrimitiveTypeTraits<T>::CppType>;
+
+    void add(const IColumn** columns, size_t /* columns_num */, size_t row_num, Arena&) {
+        const auto& vec = assert_cast<const typename PrimitiveTypeTraits<T>::ColumnType&,
+                                      TypeCheckOnRelease::DISABLE>(*columns[0])
+                                  .get_data();
+        Base::add(vec[row_num]);
+    }
+};
+
+struct AggregateFunctionEntropySingleStringData
+        : public AggregateFunctionEntropyData<UInt128, UInt128TrivialHash> {
+    using Base = AggregateFunctionEntropyData<UInt128, UInt128TrivialHash>;
+
+    void add(const IColumn** columns, size_t /* columns_num */, size_t row_num, Arena&) {
+        auto key = columns[0]->get_data_at(row_num);
+        auto hash_value = XXH_INLINE_XXH128(key.data, key.size, 0);
+        Base::add(UInt128 {hash_value.high64, hash_value.low64});
+    }
+};
+
+struct AggregateFunctionEntropyGenericData
+        : public AggregateFunctionEntropyData<UInt128, UInt128TrivialHash> {
+    using Base = AggregateFunctionEntropyData<UInt128, UInt128TrivialHash>;
+
+    void add(const IColumn** columns, size_t columns_num, size_t row_num, Arena& arena) {
+        const char* begin = nullptr;
+        StringRef key(begin, 0);
+        for (size_t i = 0; i < columns_num; ++i) {
+            auto cur_ref = columns[i]->serialize_value_into_arena(row_num, arena, begin);
+            key.data = cur_ref.data - key.size;
+            key.size += cur_ref.size;
+        }
+        auto hash_value = XXH_INLINE_XXH128(key.data, key.size, 0);
+        Base::add(UInt128 {hash_value.high64, hash_value.low64});
+    }
+};
+
+template <typename Data>
+class AggregateFunctionEntropy
+        : public IAggregateFunctionDataHelper<Data, AggregateFunctionEntropy<Data>>,
+          VarargsExpression,
+          NullableAggregateFunction {
+private:
+    size_t arguments_num;
+
+public:
+    AggregateFunctionEntropy(const DataTypes& arguments)
+            : IAggregateFunctionDataHelper<Data, AggregateFunctionEntropy<Data>>(arguments),
+              arguments_num(arguments.size()) {}
+
+    void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
+             Arena& arena) const override {
+        this->data(place).add(columns, arguments_num, row_num, arena);
+    }
+
+    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
+               Arena& 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);
+    }
+
+    void insert_result_into(ConstAggregateDataPtr place, IColumn& to) const override {
+        auto& column = assert_cast<ColumnFloat64&>(to);
+        column.get_data().push_back(this->data(place).get_result());
+    }
+
+    void reset(AggregateDataPtr place) const override { this->data(place).clear(); }
+
+    String get_name() const override { return Data::get_name(); }
+
+    DataTypePtr get_return_type() const override {
+        return std::make_shared<DataTypeNumber<TYPE_DOUBLE>>();
+    }
+};
+
+} // namespace doris::vectorized
+
+#include "common/compile_check_end.h"

be/src/vec/aggregate_functions/aggregate_function_simple_factory.cpp

@@ -79,6 +79,7 @@ void register_aggregate_function_percentile_reservoir(AggregateFunctionSimpleFac
 void register_aggregate_function_ai_agg(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_bool_union(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_sem(AggregateFunctionSimpleFactory& factory);
+void register_aggregate_function_entropy(AggregateFunctionSimpleFactory& factory);
 
 AggregateFunctionSimpleFactory& AggregateFunctionSimpleFactory::instance() {
     static std::once_flag oc;
@@ -135,6 +136,7 @@ AggregateFunctionSimpleFactory& AggregateFunctionSimpleFactory::instance() {
         register_aggregate_function_ai_agg(instance);
         register_aggregate_function_bool_union(instance);
         register_aggregate_function_sem(instance);
+        register_aggregate_function_entropy(instance);
         // Register foreach and foreachv2 functions
         register_aggregate_function_combinator_foreach(instance);
         register_aggregate_function_combinator_foreachv2(instance);

be/test/testutil/column_helper.h

@@ -127,14 +127,15 @@ struct ColumnHelper {
         return block;
     }
 
-    template <typename DataType>
-    static Block create_block(const std::vector<typename DataType::FieldType>& data1,
-                              const std::vector<typename DataType::FieldType>& data2) {
-        auto column1 = create_column<DataType>(data1);
-        auto column2 = create_column<DataType>(data2);
-        auto data_type = std::make_shared<DataType>();
-        Block block({ColumnWithTypeAndName(column1, data_type, "column1"),
-                     ColumnWithTypeAndName(column2, data_type, "column2")});
+    template <typename DataType1, typename DataType2 = DataType1>
+    static Block create_block(const std::vector<typename DataType1::FieldType>& data1,
+                              const std::vector<typename DataType2::FieldType>& data2) {
+        auto column1 = create_column<DataType1>(data1);
+        auto column2 = create_column<DataType2>(data2);
+        auto data_type1 = std::make_shared<DataType1>();
+        auto data_type2 = std::make_shared<DataType2>();
+        Block block({ColumnWithTypeAndName(column1, data_type1, "column1"),
+                     ColumnWithTypeAndName(column2, data_type2, "column2")});
         return block;
     }