[Feature] Add Topn udaf

be/src/common/daemon.cpp

@@ -41,6 +41,7 @@
 #include "exprs/string_functions.h"
 #include "exprs/time_operators.h"
 #include "exprs/timestamp_functions.h"
+#include "exprs/topn_function.h"
 #include "exprs/utility_functions.h"
 #include "geo/geo_functions.h"
 #include "olap/options.h"
@@ -261,6 +262,7 @@ void Daemon::init(int argc, char** argv, const std::vector<StorePath>& paths) {
     BitmapFunctions::init();
     HllFunctions::init();
     HashFunctions::init();
+    TopNFunctions::init();
 
     LOG(INFO) << CpuInfo::debug_string();
     LOG(INFO) << DiskInfo::debug_string();

be/src/exprs/CMakeLists.txt

@@ -65,4 +65,5 @@ add_library(Exprs
   new_agg_fn_evaluator.cc
   bitmap_function.cpp
   hll_function.cpp
-  grouping_sets_functions.cpp)
+  grouping_sets_functions.cpp
+  topn_function.cpp)

be/src/exprs/topn_function.cpp

@@ -0,0 +1,119 @@
+// 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 "exprs/topn_function.h"
+#include "util/topn_counter.h"
+#include "util/slice.h"
+
+namespace doris {
+
+using doris_udf::AnyVal;
+
+void TopNFunctions::init() {
+}
+
+void TopNFunctions::topn_init(FunctionContext* ctx, StringVal *dst) {
+    dst->is_null = false;
+    dst->len = sizeof(TopNCounter);
+    const AnyVal* space_expand_rate_val = ctx->get_constant_arg(2);
+    if (space_expand_rate_val != nullptr) {
+        int32_t space_expand_rate = reinterpret_cast<const IntVal*>(space_expand_rate_val)->val;
+        dst->ptr = (uint8_t *) new TopNCounter(space_expand_rate);
+        return;
+    }
+    dst->ptr = (uint8_t *) new TopNCounter();
+}
+
+template <typename T>
+void TopNFunctions::topn_update(FunctionContext*, const T& src, const IntVal& topn, StringVal* dst) {
+    if (src.is_null) {
+        return;
+    }
+    auto* dst_topn = reinterpret_cast<TopNCounter*>(dst->ptr);
+    dst_topn->set_top_num(topn.val);
+    dst_topn->add_item(src);
+}
+
+template <typename T>
+void TopNFunctions::topn_update(FunctionContext *, const T &src, const IntVal &topn, const IntVal &space_expand_rate,
+                                StringVal *dst) {
+    if (src.is_null) {
+        return;
+    }
+    auto* dst_topn = reinterpret_cast<TopNCounter*>(dst->ptr);
+    dst_topn->set_top_num(topn.val);
+    dst_topn->add_item(src);
+}
+
+void TopNFunctions::topn_merge(FunctionContext* ctx, const StringVal &src, StringVal *dst) {
+    if (src.is_null) {
+        return;
+    }
+    auto* dst_topn = reinterpret_cast<TopNCounter*>(dst->ptr);
+    dst_topn->merge(TopNCounter(Slice(src.ptr, src.len)));
+}
+
+StringVal TopNFunctions::topn_serialize(FunctionContext *ctx, const StringVal &src) {
+    auto* src_topn = reinterpret_cast<TopNCounter*>(src.ptr);
+
+    std::string buffer;
+    src_topn->serialize(&buffer);
+    StringVal result(ctx, buffer.size());
+    memcpy(result.ptr, buffer.data(), buffer.size());
+    delete src_topn;
+    return result;
+}
+
+StringVal TopNFunctions::topn_finalize(FunctionContext* ctx, const StringVal &src) {
+    auto* src_topn = reinterpret_cast<TopNCounter*>(src.ptr);
+    std::string result_str;
+    src_topn->finalize(result_str);
+
+    StringVal result(ctx, result_str.size());
+    memcpy(result.ptr, result_str.data(), result_str.size());
+
+    delete src_topn;
+    return result;
+}
+
+template void TopNFunctions::topn_update(FunctionContext*, const BooleanVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const TinyIntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const SmallIntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const BigIntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const FloatVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const DoubleVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const StringVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DateTimeVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const LargeIntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DecimalVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DecimalV2Val&, const IntVal&, StringVal*);
+
+template void TopNFunctions::topn_update(FunctionContext*, const BooleanVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const TinyIntVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const SmallIntVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const IntVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const BigIntVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const FloatVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const DoubleVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext *, const StringVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DateTimeVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const LargeIntVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DecimalVal&, const IntVal&, const IntVal&, StringVal*);
+template void TopNFunctions::topn_update(FunctionContext*, const DecimalV2Val&, const IntVal&, const IntVal&, StringVal*);
+
+}

be/src/exprs/topn_function.h

@@ -0,0 +1,47 @@
+// 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.
+
+#ifndef DORIS_BE_SRC_EXPRS_TOPN_FUNCTION_H
+#define DORIS_BE_SRC_EXPRS_TOPN_FUNCTION_H
+
+#include "udf/udf.h"
+
+namespace doris {
+
+class TopNFunctions {
+public:
+    static void init();
+
+    static void topn_init(FunctionContext*, StringVal* dst);
+
+    template <typename T>
+    static void topn_update(FunctionContext*, const T& src, const IntVal& topn, StringVal* dst);
+
+    template <typename T>
+    static void topn_update(FunctionContext*, const T& src, const IntVal& topn, const IntVal& space_expand_rate,
+            StringVal* dst);
+
+    static void topn_merge(FunctionContext*,const StringVal& src, StringVal* dst);
+
+    static StringVal topn_serialize(FunctionContext* ctx, const StringVal& src);
+
+    static StringVal topn_finalize(FunctionContext*, const StringVal& src);
+};
+
+}
+
+#endif //DORIS_BE_SRC_EXPRS_TOPN_FUNCTION_H

be/src/util/CMakeLists.txt

@@ -102,6 +102,7 @@ set(UTIL_FILES
   brpc_stub_cache.cpp
   zlib.cpp
   pprof_utils.cpp
+  topn_counter.cpp
 )
 
 if (WITH_MYSQL)

be/src/util/topn_counter.cpp

@@ -0,0 +1,142 @@
+// 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 <algorithm>
+#include <rapidjson/writer.h>
+#include <rapidjson/stringbuffer.h>
+
+#include "gen_cpp/olap_common.pb.h"
+#include "topn_counter.h"
+#include "slice.h"
+
+namespace doris {
+
+void TopNCounter::add_item(const std::string& item, uint64_t incrementCount) {
+    auto iter = _counter_map->find(item);
+    if (iter != _counter_map->end()) {
+        iter->second.add_count(incrementCount);
+    } else {
+        _counter_map->insert(std::make_pair(item, Counter(item, incrementCount)));
+    }
+    _ordered = false;
+}
+
+void TopNCounter::serialize(std::string* buffer) {
+    sort_retain(_capacity);
+    PTopNCounter topn_counter;
+    topn_counter.set_top_num(_top_num);
+    topn_counter.set_space_expand_rate(_space_expand_rate);
+    for(std::vector<Counter>::const_iterator it = _counter_vec->begin(); it != _counter_vec->end(); ++it)
+    {
+        PCounter* counter = topn_counter.add_counter();
+        counter->set_item(it->get_item());
+        counter->set_count(it->get_count());
+    }
+    topn_counter.SerializeToString(buffer);
+}
+
+bool TopNCounter::deserialize(const doris::Slice &src) {
+    PTopNCounter topn_counter;
+    if (!topn_counter.ParseFromArray(src.data, src.size)) {
+        LOG(WARNING) << "topn counter deserialize failed";
+        return false;
+    }
+
+    _space_expand_rate = topn_counter.space_expand_rate();
+    set_top_num(topn_counter.top_num());
+    for (int i = 0; i < topn_counter.counter_size(); ++i) {
+        const PCounter& counter = topn_counter.counter(i);
+        _counter_map->insert(std::make_pair(counter.item(), Counter(counter.item(), counter.count())));
+        _counter_vec->emplace_back(counter.item(), counter.count());
+    }
+    _ordered = true;
+    return true;
+}
+
+void TopNCounter::sort_retain(uint32_t capacity) {
+    _counter_vec->clear();
+    sort_retain(capacity, _counter_vec);
+    _ordered = true;
+}
+
+void TopNCounter::sort_retain(uint32_t capacity, std::vector<Counter>* sort_vec) {
+    for(std::unordered_map<std::string, Counter>::const_iterator it = _counter_map->begin(); it != _counter_map->end(); ++it) {
+        sort_vec->emplace_back(it->second.get_item(), it->second.get_count());
+    }
+
+    std::sort(sort_vec->begin(), sort_vec->end(), TopNComparator());
+    if (sort_vec->size() > capacity) {
+        for (uint32_t i = 0, n = sort_vec->size() - capacity; i < n; ++i) {
+            auto &counter = sort_vec->back();
+            _counter_map->erase(counter.get_item());
+            sort_vec->pop_back();
+        }
+    }
+}
+
+// Based on the  parallel version of the Space Saving algorithm as described in:
+// A parallel space saving algorithm for frequent items and the Hurwitz zeta distribution by Massimo Cafaro, et al.
+void TopNCounter::merge(doris::TopNCounter &&other) {
+    if (other._counter_map->size() == 0) {
+        return;
+    }
+
+    _space_expand_rate = other._space_expand_rate;
+    set_top_num(other._top_num);
+    bool this_full = _counter_map->size() >= _capacity;
+    bool another_full = other._counter_map->size() >= other._capacity;
+
+    uint64_t m1 = this_full ? _counter_vec->back().get_count() : 0;
+    uint64_t m2 = another_full ? other._counter_vec->back().get_count() : 0;
+
+    if (another_full == true) {
+        for (auto &entry : *(this->_counter_map)) {
+            entry.second.add_count(m2);
+        }
+    }
+
+    for (auto &other_entry : *(other._counter_map)) {
+        auto itr = this->_counter_map->find(other_entry.first);
+        if (itr != _counter_map->end()) {
+            itr->second.add_count(other_entry.second.get_count() - m2);
+        } else {
+            this->_counter_map->insert(std::make_pair(other_entry.first,
+                    Counter(other_entry.first,other_entry.second.get_count() + m1)));
+        }
+    }
+    _ordered = false;
+    sort_retain(_capacity);
+}
+
+void TopNCounter::finalize(std::string& finalize_str) {
+    if (!_ordered) {
+        sort_retain(_top_num);
+    }
+    // use json format print
+    rapidjson::StringBuffer buffer;
+    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+    uint32_t k = 0;
+    writer.StartObject();
+    for (std::vector<Counter>::const_iterator it = _counter_vec->begin(); it != _counter_vec->end() && k < _top_num; ++it, ++k) {
+        writer.Key(it->get_item().data());
+        writer.Uint64(it->get_count());
+    }
+    writer.EndObject();
+    finalize_str = buffer.GetString();
+}
+
+}