[Feature] Support exact percentile aggregate function

be/src/exprs/aggregate_functions.cpp

@@ -31,6 +31,7 @@
 #include "runtime/string_value.h"
 #include "util/debug_util.h"
 #include "util/tdigest.h"
+#include "util/counts.h"
 
 // TODO: this file should be cross compiled and then all of the builtin
 // aggregate functions will have a codegen enabled path. Then we can remove
@@ -172,6 +173,76 @@ void AggregateFunctions::count_remove(FunctionContext*, const AnyVal& src, BigIn
     }
 }
 
+struct PercentileState {
+    Counts counts; 
+    double quantile = -1.0;
+};
+
+void AggregateFunctions::percentile_init(FunctionContext* ctx, StringVal* dst) {
+    dst->is_null = false;
+    dst->len = sizeof(PercentileState);
+    dst->ptr = (uint8_t*) new PercentileState();
+}
+
+template <typename T>
+void AggregateFunctions::percentile_update(FunctionContext* ctx, const T& src,
+                                            const DoubleVal& quantile, StringVal* dst) {
+    if (src.is_null) {
+        return;
+    }
+
+    DCHECK(dst->ptr != nullptr);
+    DCHECK_EQ(sizeof(PercentileState), dst->len);
+
+    PercentileState* percentile = reinterpret_cast<PercentileState*>(dst->ptr);
+    percentile->counts.increment(src.val, 1);
+    percentile->quantile = quantile.val;
+}
+
+void AggregateFunctions::percentile_merge(FunctionContext* ctx, const StringVal& src, StringVal* dst) {
+    DCHECK(dst->ptr != nullptr);
+    DCHECK_EQ(sizeof(PercentileState), dst->len);
+
+    double quantile;
+    memcpy(&quantile, src.ptr, sizeof(double));
+
+    PercentileState* src_percentile = new PercentileState();
+    src_percentile->quantile = quantile;
+    src_percentile->counts.unserialize(src.ptr + sizeof(double));
+
+    PercentileState* dst_percentile = reinterpret_cast<PercentileState*>(dst->ptr);
+    dst_percentile->counts.merge(&src_percentile->counts);
+    if (dst_percentile->quantile == -1.0) {
+        dst_percentile->quantile = quantile;
+    }
+
+    delete src_percentile;
+}
+
+StringVal AggregateFunctions::percentile_serialize(FunctionContext* ctx, const StringVal& src) {
+    DCHECK(!src.is_null);
+
+    PercentileState* percentile = reinterpret_cast<PercentileState*>(src.ptr);
+    uint32_t serialize_size = percentile->counts.serialized_size();
+    StringVal result(ctx, sizeof(double) + serialize_size);
+    memcpy(result.ptr, &percentile->quantile, sizeof(double));
+    percentile->counts.serialize(result.ptr + sizeof(double));
+
+    delete percentile;
+    return result;
+}
+
+DoubleVal AggregateFunctions::percentile_finalize(FunctionContext* ctx, const StringVal& src) {
+    DCHECK(!src.is_null);
+
+    PercentileState* percentile = reinterpret_cast<PercentileState*>(src.ptr);
+    double quantile = percentile->quantile;
+    auto result = percentile->counts.terminate(quantile);
+
+    delete percentile;
+    return result;
+}
+
 struct PercentileApproxState {
 public:
     PercentileApproxState() : digest(new TDigest()) {}
@@ -2696,6 +2767,9 @@ template void AggregateFunctions::offset_fn_update<DecimalV2Val>(FunctionContext
                                                                  const DecimalV2Val&,
                                                                  DecimalV2Val* dst);
 
+template void AggregateFunctions::percentile_update<BigIntVal>(
+        FunctionContext* ctx, const BigIntVal&, const DoubleVal&, StringVal*);
+
 template void AggregateFunctions::percentile_approx_update<doris_udf::DoubleVal>(
         FunctionContext* ctx, const doris_udf::DoubleVal&, const doris_udf::DoubleVal&,
         doris_udf::StringVal*);

be/src/exprs/aggregate_functions.h

@@ -74,6 +74,19 @@ class AggregateFunctions {
 
     static void count_star_remove(FunctionContext*, BigIntVal* dst);
 
+    // Implementation of percentile
+    static void percentile_init(FunctionContext* ctx, StringVal* dst);
+
+    template <typename T>
+    static void percentile_update(FunctionContext* ctx, const T& src,
+                                    const DoubleVal& quantile, StringVal* dst);
+
+    static void percentile_merge(FunctionContext* ctx, const StringVal& src, StringVal* dst);
+
+    static StringVal percentile_serialize(FunctionContext* ctx, const StringVal& state_sv);
+
+    static DoubleVal percentile_finalize(FunctionContext* ctx, const StringVal& src);
+
     // Implementation of percentile_approx
     static void percentile_approx_init(doris_udf::FunctionContext* ctx, doris_udf::StringVal* dst);
 

be/src/util/counts.h

@@ -0,0 +1,138 @@
+// 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_UTIL_COUNTS_H_
+#define DORIS_BE_SRC_UTIL_COUNTS_H_
+
+#include <algorithm>
+#include <cmath>
+#include <unordered_map>
+#include <vector>
+
+#include "udf/udf.h"
+
+namespace doris {
+
+class Counts {
+public:
+    Counts() = default;
+
+    inline void merge(const Counts* other) {
+        if (other == nullptr || other->_counts.empty()) {
+            return;
+        }
+
+        for (auto& cell : other->_counts) {
+            increment(cell.first, cell.second);
+        }
+    }
+
+    void increment(int64_t key, uint32_t i) {
+        auto item = _counts.find(key);
+        if (item != _counts.end()) {
+            item->second += i;
+        } else {
+            _counts.emplace(std::make_pair(key, i));
+        }
+    }
+
+    uint32_t serialized_size() {
+        return sizeof(uint32_t) + sizeof(int64_t) * _counts.size() +
+               sizeof(uint32_t) * _counts.size();
+    }
+
+    void serialize(uint8_t* writer) {
+        uint32_t size = _counts.size();
+        memcpy(writer, &size, sizeof(uint32_t));
+        writer += sizeof(uint32_t);
+        for (auto& cell : _counts) {
+            memcpy(writer, &cell.first, sizeof(int64_t));
+            writer += sizeof(int64_t);
+            memcpy(writer, &cell.second, sizeof(uint32_t));
+            writer += sizeof(uint32_t);
+        }
+    }
+
+    void unserialize(const uint8_t* type_reader) {
+        uint32_t size;
+        memcpy(&size, type_reader, sizeof(uint32_t));
+        type_reader += sizeof(uint32_t);
+        for (uint32_t i = 0; i < size; ++i) {
+            int64_t key;
+            uint32_t count;
+            memcpy(&key, type_reader, sizeof(int64_t));
+            type_reader += sizeof(int64_t);
+            memcpy(&count, type_reader, sizeof(uint32_t));
+            type_reader += sizeof(uint32_t);
+            _counts.emplace(std::make_pair(key, count));
+        }
+    }
+
+    double get_percentile(std::vector<std::pair<int64_t, uint32_t>>& counts, double position) {
+        long lower = std::floor(position);
+        long higher = std::ceil(position);
+
+        auto iter = counts.begin();
+        for (; iter != counts.end() && iter->second < lower + 1; ++iter)
+            ;
+
+        int64_t lower_key = iter->first;
+        if (higher == lower) {
+            return lower_key;
+        }
+
+        if (iter->second < higher + 1) {
+            iter++;
+        }
+
+        int64_t higher_key = iter->first;
+        if (lower_key == higher_key) {
+            return lower_key;
+        }
+
+        return (higher - position) * lower_key + (position - lower) * higher_key;
+    }
+
+    doris_udf::DoubleVal terminate(double quantile) {
+        if (_counts.empty()) {
+            return doris_udf::DoubleVal();
+        }
+
+        std::vector<std::pair<int64_t, uint32_t>> elems(_counts.begin(), _counts.end());
+        sort(elems.begin(), elems.end(),
+             [](const std::pair<int64_t, uint32_t> l, const std::pair<int64_t, uint32_t> r) {
+                 return l.first < r.first;
+             });
+
+        long total = 0;
+        for (auto& cell : elems) {
+            total += cell.second;
+            cell.second = total;
+        }
+
+        long max_position = total - 1;
+        double position = max_position * quantile;
+        return doris_udf::DoubleVal(get_percentile(elems, position));
+    }
+
+private:
+    std::unordered_map<int64_t, uint32_t> _counts;
+};
+
+} // namespace doris
+
+#endif // DORIS_BE_SRC_UTIL_COUNTS_H_

be/test/exprs/CMakeLists.txt

@@ -29,6 +29,7 @@ ADD_BE_TEST(json_function_test)
 ADD_BE_TEST(string_functions_test)
 ADD_BE_TEST(timestamp_functions_test)
 ADD_BE_TEST(percentile_approx_test)
+ADD_BE_TEST(percentile_test)
 ADD_BE_TEST(bitmap_function_test)
 ADD_BE_TEST(hll_function_test)
 ADD_BE_TEST(encryption_functions_test)

be/test/exprs/percentile_test.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 <gtest/gtest.h>
+
+#include "exprs/aggregate_functions.h"
+#include "testutil/function_utils.h"
+
+namespace doris {
+
+class PercentileTest : public testing::Test {
+public:
+    PercentileTest() {}
+};
+
+TEST_F(PercentileTest, testSample) {
+    FunctionUtils* futil = new FunctionUtils();
+    doris_udf::FunctionContext* context = futil->get_fn_ctx();
+
+    DoubleVal doubleQ(0.9);
+
+    StringVal stringVal1;
+    BigIntVal int1(1);
+    AggregateFunctions::percentile_init(context, &stringVal1);
+    AggregateFunctions::percentile_update(context, int1, doubleQ, &stringVal1);
+    BigIntVal int2(2);
+    AggregateFunctions::percentile_update(context, int2, doubleQ, &stringVal1);
+
+    StringVal s = AggregateFunctions::percentile_serialize(context, stringVal1);
+
+    StringVal stringVal2;
+    AggregateFunctions::percentile_init(context, &stringVal2);
+    AggregateFunctions::percentile_merge(context, s, &stringVal2);
+    DoubleVal v = AggregateFunctions::percentile_finalize(context, stringVal2);
+    ASSERT_EQ(v.val, 1.9);
+    delete futil;
+}
+
+TEST_F(PercentileTest, testNoMerge) {
+    FunctionUtils* futil = new FunctionUtils();
+    doris_udf::FunctionContext* context = futil->get_fn_ctx();
+
+    DoubleVal doubleQ(0.9);
+
+    StringVal stringVal1;
+    BigIntVal val(1);
+    AggregateFunctions::percentile_init(context, &stringVal1);
+    AggregateFunctions::percentile_update(context, val, doubleQ, &stringVal1);
+    BigIntVal val2(2);
+    AggregateFunctions::percentile_update(context, val2, doubleQ, &stringVal1);
+
+    DoubleVal v = AggregateFunctions::percentile_finalize(context, stringVal1);
+    ASSERT_EQ(v.val, 1.9);
+    delete futil;
+}
+
+TEST_F(PercentileTest, testSerialize) {
+    FunctionUtils* futil = new FunctionUtils();
+    doris_udf::FunctionContext* context = futil->get_fn_ctx();
+
+    DoubleVal doubleQ(0.999);
+    StringVal stringVal;
+    AggregateFunctions::percentile_init(context, &stringVal);
+
+    for (int i = 1; i <= 100000; i++) {
+        BigIntVal val(i);
+        AggregateFunctions::percentile_update(context, val, doubleQ, &stringVal);
+    }
+    StringVal serialized = AggregateFunctions::percentile_serialize(context, stringVal);
+
+    // mock serialize
+    StringVal stringVal2;
+    AggregateFunctions::percentile_init(context, &stringVal2);
+    AggregateFunctions::percentile_merge(context, serialized, &stringVal2);
+    DoubleVal v = AggregateFunctions::percentile_finalize(context, stringVal2);
+    ASSERT_DOUBLE_EQ(v.val, 99900.001);
+
+    // merge init percentile stringVal3 should not change the correct result
+    AggregateFunctions::percentile_init(context, &stringVal);
+
+    for (int i = 1; i <= 100000; i++) {
+        BigIntVal val(i);
+        AggregateFunctions::percentile_update(context, val, doubleQ, &stringVal);
+    }
+    serialized = AggregateFunctions::percentile_serialize(context, stringVal);
+
+    StringVal stringVal3;
+    AggregateFunctions::percentile_init(context, &stringVal2);
+    AggregateFunctions::percentile_init(context, &stringVal3);
+    StringVal serialized2 = AggregateFunctions::percentile_serialize(context, stringVal3);
+
+    AggregateFunctions::percentile_merge(context, serialized, &stringVal2);
+    AggregateFunctions::percentile_merge(context, serialized2, &stringVal2);
+    v = AggregateFunctions::percentile_finalize(context, stringVal2);
+    ASSERT_DOUBLE_EQ(v.val, 99900.001);
+
+    delete futil;
+}
+
+} // namespace doris
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}

be/test/util/CMakeLists.txt

@@ -72,3 +72,4 @@ ADD_BE_TEST(s3_uri_test)
 ADD_BE_TEST(s3_storage_backend_test)
 ADD_BE_TEST(broker_storage_backend_test)
 ADD_BE_TEST(sort_heap_test)
+ADD_BE_TEST(counts_test)