[Vectorized][Feature] Support min_by/max_by function.

be/src/vec/CMakeLists.txt

@@ -26,6 +26,7 @@ set(VEC_FILES
   aggregate_functions/aggregate_function_distinct.cpp
   aggregate_functions/aggregate_function_sum.cpp
   aggregate_functions/aggregate_function_min_max.cpp
+  aggregate_functions/aggregate_function_min_max_by.cpp
   aggregate_functions/aggregate_function_null.cpp
   aggregate_functions/aggregate_function_uniq.cpp
   aggregate_functions/aggregate_function_hll_union_agg.cpp

be/src/vec/aggregate_functions/aggregate_function_min_max_by.cpp

@@ -0,0 +1,116 @@
+// 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 "vec/aggregate_functions/aggregate_function_min_max.h"
+#include "vec/aggregate_functions/aggregate_function_min_max_by.h"
+
+#include "vec/aggregate_functions/aggregate_function_simple_factory.h"
+#include "vec/aggregate_functions/factory_helpers.h"
+#include "vec/aggregate_functions/helpers.h"
+
+namespace doris::vectorized {
+
+/// min_by, max_by
+template <template <typename, bool> class AggregateFunctionTemplate,
+          template <typename, typename> class Data, typename VT>
+static IAggregateFunction* create_aggregate_function_min_max_by_impl(
+        const DataTypes& argument_types) {
+    const DataTypePtr& value_arg_type = argument_types[0];
+    const DataTypePtr& key_arg_type = argument_types[1];
+
+    WhichDataType which(key_arg_type);
+#define DISPATCH(TYPE)                                                                     \
+    if (which.idx == TypeIndex::TYPE)                                                      \
+        return new AggregateFunctionTemplate<Data<VT, SingleValueDataFixed<TYPE>>, false>( \
+                value_arg_type, key_arg_type);
+    FOR_NUMERIC_TYPES(DISPATCH)
+#undef DISPATCH
+    if (which.idx == TypeIndex::String) {
+        return new AggregateFunctionTemplate<Data<VT, SingleValueDataString>, false>(value_arg_type,
+                                                                                     key_arg_type);
+    }
+    if (which.idx == TypeIndex::DateTime || which.idx == TypeIndex::Date) {
+        return new AggregateFunctionTemplate<Data<VT, SingleValueDataFixed<Int64>>, false>(
+                value_arg_type, key_arg_type);
+    }
+    if (which.idx == TypeIndex::Decimal128) {
+        return new AggregateFunctionTemplate<Data<VT, SingleValueDataFixed<DecimalV2Value>>, false>(
+                value_arg_type, key_arg_type);
+    }
+    return nullptr;
+}
+
+/// min_by, max_by
+template <template <typename, bool> class AggregateFunctionTemplate,
+          template <typename, typename> class Data>
+static IAggregateFunction* create_aggregate_function_min_max_by(const String& name,
+                                                                const DataTypes& argument_types,
+                                                                const Array& parameters) {
+    assert_no_parameters(name, parameters);
+    assert_binary(name, argument_types);
+
+    const DataTypePtr& value_arg_type = argument_types[0];
+
+    WhichDataType which(value_arg_type);
+#define DISPATCH(TYPE)                                                                    \
+    if (which.idx == TypeIndex::TYPE)                                                     \
+        return create_aggregate_function_min_max_by_impl<AggregateFunctionTemplate, Data, \
+                                                         SingleValueDataFixed<TYPE>>(     \
+                argument_types);
+    FOR_NUMERIC_TYPES(DISPATCH)
+#undef DISPATCH
+    if (which.idx == TypeIndex::String) {
+        return create_aggregate_function_min_max_by_impl<AggregateFunctionTemplate, Data,
+                                                         SingleValueDataString>(argument_types);
+    }
+    if (which.idx == TypeIndex::DateTime || which.idx == TypeIndex::Date) {
+        return create_aggregate_function_min_max_by_impl<AggregateFunctionTemplate, Data,
+                                                         SingleValueDataFixed<Int64>>(
+                argument_types);
+    }
+    if (which.idx == TypeIndex::Decimal128) {
+        return create_aggregate_function_min_max_by_impl<AggregateFunctionTemplate, Data,
+                                                         SingleValueDataFixed<DecimalV2Value>>(
+                argument_types);
+    }
+    return nullptr;
+}
+
+AggregateFunctionPtr create_aggregate_function_max_by(const std::string& name,
+                                                      const DataTypes& argument_types,
+                                                      const Array& parameters,
+                                                      const bool result_is_nullable) {
+    return AggregateFunctionPtr(create_aggregate_function_min_max_by<AggregateFunctionsMinMaxBy,
+                                                                     AggregateFunctionMaxByData>(
+            name, argument_types, parameters));
+}
+
+AggregateFunctionPtr create_aggregate_function_min_by(const std::string& name,
+                                                      const DataTypes& argument_types,
+                                                      const Array& parameters,
+                                                      const bool result_is_nullable) {
+    return AggregateFunctionPtr(create_aggregate_function_min_max_by<AggregateFunctionsMinMaxBy,
+                                                                     AggregateFunctionMinByData>(
+            name, argument_types, parameters));
+}
+
+void register_aggregate_function_min_max_by(AggregateFunctionSimpleFactory& factory) {
+    factory.register_function("max_by", create_aggregate_function_max_by);
+    factory.register_function("min_by", create_aggregate_function_min_by);
+}
+
+} // namespace doris::vectorized

be/src/vec/aggregate_functions/aggregate_function_min_max_by.h

@@ -0,0 +1,161 @@
+// 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 "common/logging.h"
+#include "vec/aggregate_functions/aggregate_function.h"
+#include "vec/columns/column_decimal.h"
+#include "vec/columns/column_vector.h"
+#include "vec/common/assert_cast.h"
+#include "vec/io/io_helper.h"
+
+namespace doris::vectorized {
+template <typename VT, typename KT>
+struct AggregateFunctionMinMaxByBaseData {
+protected:
+    VT value;
+    KT key;
+
+public:
+    void insert_result_into(IColumn& to) const { value.insert_result_into(to); }
+
+    void reset() {
+        value.reset();
+        key.reset();
+    }
+    void write(BufferWritable& buf) const {
+        value.write(buf);
+        key.write(buf);
+    }
+
+    void read(BufferReadable& buf) {
+        value.read(buf);
+        key.read(buf);
+    }
+};
+
+template <typename VT, typename KT>
+struct AggregateFunctionMaxByData : public AggregateFunctionMinMaxByBaseData<VT, KT> {
+    using Self = AggregateFunctionMaxByData;
+    bool change_if_better(const IColumn& value_column, const IColumn& key_column, size_t row_num,
+                          Arena* arena) {
+        if (this->key.change_if_greater(key_column, row_num, arena)) {
+            this->value.change(value_column, row_num, arena);
+            return true;
+        }
+        return false;
+    }
+
+    bool change_if_better(const Self& to, Arena* arena) {
+        if (this->key.change_if_greater(to.key, arena)) {
+            this->value.change(to.value, arena);
+            return true;
+        }
+        return false;
+    }
+
+    static const char* name() { return "max_by"; }
+};
+
+template <typename VT, typename KT>
+struct AggregateFunctionMinByData : public AggregateFunctionMinMaxByBaseData<VT, KT> {
+    using Self = AggregateFunctionMinByData;
+    bool change_if_better(const IColumn& value_column, const IColumn& key_column, size_t row_num,
+                          Arena* arena) {
+        if (this->key.change_if_less(key_column, row_num, arena)) {
+            this->value.change(value_column, row_num, arena);
+            return true;
+        }
+        return false;
+    }
+
+    bool change_if_better(const Self& to, Arena* arena) {
+        if (this->key.change_if_less(to.key, arena)) {
+            this->value.change(to.value, arena);
+            return true;
+        }
+        return false;
+    }
+
+    static const char* name() { return "min_by"; }
+};
+
+template <typename Data, bool AllocatesMemoryInArena>
+class AggregateFunctionsMinMaxBy final
+        : public IAggregateFunctionDataHelper<
+                  Data, AggregateFunctionsMinMaxBy<Data, AllocatesMemoryInArena>> {
+private:
+    DataTypePtr& value_type;
+    DataTypePtr& key_type;
+
+public:
+    AggregateFunctionsMinMaxBy(const DataTypePtr& value_type_, const DataTypePtr& key_type_)
+            : IAggregateFunctionDataHelper<
+                      Data, AggregateFunctionsMinMaxBy<Data, AllocatesMemoryInArena>>(
+                      {value_type_, key_type_}, {}),
+              value_type(this->argument_types[0]),
+              key_type(this->argument_types[1]) {
+        if (StringRef(Data::name()) == StringRef("min_by") ||
+            StringRef(Data::name()) == StringRef("max_by")) {
+            CHECK(key_type_->is_comparable());
+        }
+    }
+
+    String get_name() const override { return Data::name(); }
+
+    DataTypePtr get_return_type() const override { return value_type; }
+
+    void add(AggregateDataPtr __restrict place, const IColumn** columns, size_t row_num,
+             Arena* arena) const override {
+        this->data(place).change_if_better(*columns[0], *columns[1], row_num, arena);
+    }
+
+    void reset(AggregateDataPtr place) const override { this->data(place).reset(); }
+
+    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
+               Arena* arena) const override {
+        this->data(place).change_if_better(this->data(rhs), arena);
+    }
+
+    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);
+    }
+
+    bool allocates_memory_in_arena() const override { return AllocatesMemoryInArena; }
+
+    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {
+        this->data(place).insert_result_into(to);
+    }
+};
+
+AggregateFunctionPtr create_aggregate_function_max_by(const std::string& name,
+                                                      const DataTypes& argument_types,
+                                                      const Array& parameters,
+                                                      const bool result_is_nullable);
+
+AggregateFunctionPtr create_aggregate_function_min_by(const std::string& name,
+                                                      const DataTypes& argument_types,
+                                                      const Array& parameters,
+                                                      const bool result_is_nullable);
+
+} // namespace doris::vectorized

be/src/vec/aggregate_functions/aggregate_function_simple_factory.cpp

@@ -29,6 +29,7 @@ class AggregateFunctionSimpleFactory;
 void register_aggregate_function_sum(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_combinator_null(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_minmax(AggregateFunctionSimpleFactory& factory);
+void register_aggregate_function_min_max_by(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_avg(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_count(AggregateFunctionSimpleFactory& factory);
 void register_aggregate_function_HLL_union_agg(AggregateFunctionSimpleFactory& factory);
@@ -51,6 +52,7 @@ AggregateFunctionSimpleFactory& AggregateFunctionSimpleFactory::instance() {
     std::call_once(oc, [&]() {
         register_aggregate_function_sum(instance);
         register_aggregate_function_minmax(instance);
+        register_aggregate_function_min_max_by(instance);
         register_aggregate_function_avg(instance);
         register_aggregate_function_count(instance);
         register_aggregate_function_uniq(instance);

be/test/CMakeLists.txt

@@ -331,6 +331,7 @@ set(VEC_TEST_FILES
     vec/aggregate_functions/agg_test.cpp
     vec/aggregate_functions/agg_min_max_test.cpp
     vec/aggregate_functions/vec_window_funnel_test.cpp
+    vec/aggregate_functions/agg_min_max_by_test.cpp
     vec/core/block_test.cpp
     vec/core/column_array_test.cpp
     vec/core/column_complex_test.cpp