From 6c064d62e92496c5eba383261576233552156585 Mon Sep 17 00:00:00 2001
From: feng lv <fenglv15@mails.ucas.ac.cn>
Date: Wed, 12 May 2021 12:29:43 +0000
Subject: [PATCH] add rangeSum udaf

fix
---
 .../AggregateFunctionRangeSum.cpp             |  56 +++++
 .../AggregateFunctionRangeSum.h               | 194 ++++++++++++++++++
 src/AggregateFunctions/Helpers.h              |  11 +
 .../registerAggregateFunctions.cpp            |   3 +
 4 files changed, 264 insertions(+)
 create mode 100644 src/AggregateFunctions/AggregateFunctionRangeSum.cpp
 create mode 100644 src/AggregateFunctions/AggregateFunctionRangeSum.h

diff --git a/src/AggregateFunctions/AggregateFunctionRangeSum.cpp b/src/AggregateFunctions/AggregateFunctionRangeSum.cpp
new file mode 100644
index 00000000000..4584eea7a39
--- /dev/null
+++ b/src/AggregateFunctions/AggregateFunctionRangeSum.cpp
@@ -0,0 +1,56 @@
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <AggregateFunctions/AggregateFunctionRangeSum.h>
+#include <AggregateFunctions/FactoryHelpers.h>
+#include <AggregateFunctions/Helpers.h>
+#include <DataTypes/DataTypeDate.h>
+#include <DataTypes/DataTypeDateTime.h>
+
+#include <ext/range.h>
+
+
+namespace DB
+{
+namespace ErrorCodes
+{
+    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
+}
+
+namespace
+{
+    template <template <typename> class Data>
+    AggregateFunctionPtr createAggregateFunctionRangeSum(const std::string & name, const DataTypes & arguments, const Array &)
+    {
+        if (arguments.size() != 2)
+            throw Exception(
+                "Aggregate function " + name + " requires two timestamps argument.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+        for (const auto i : ext::range(1, arguments.size()))
+        {
+            const auto * arg = arguments[i].get();
+            if (!WhichDataType(arg).isNativeUInt())
+                throw Exception{
+                    "Illegal type " + arg->getName() + " of argument " + toString(i + 1) + " of aggregate function " + name
+                        + ", must be Native Unsigned Number",
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+        }
+
+        AggregateFunctionPtr res(createWithNativeUnsignedIntegerType<AggregateFunctionRangeSum, Data>(*arguments[0], arguments));
+
+        if (res)
+            return res;
+
+        throw Exception{
+            "Illegal type " + arguments.front().get()->getName() + " of first argument of aggregate function " + name
+                + ", must be Native Unsigned Number",
+            ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+}
+
+}
+
+void registerAggregateFunctionRangeSum(AggregateFunctionFactory & factory)
+{
+    factory.registerFunction("rangeSum", createAggregateFunctionRangeSum<AggregateFunctionRangeSumData>);
+}
+
+}
diff --git a/src/AggregateFunctions/AggregateFunctionRangeSum.h b/src/AggregateFunctions/AggregateFunctionRangeSum.h
new file mode 100644
index 00000000000..66968bf6a3e
--- /dev/null
+++ b/src/AggregateFunctions/AggregateFunctionRangeSum.h
@@ -0,0 +1,194 @@
+#pragma once
+
+#include <unordered_set>
+#include <Columns/ColumnsNumber.h>
+#include <DataTypes/DataTypeDateTime.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <IO/ReadHelpers.h>
+#include <IO/WriteHelpers.h>
+#include <Common/ArenaAllocator.h>
+#include <Common/assert_cast.h>
+
+#include <AggregateFunctions/AggregateFunctionNull.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int BAD_ARGUMENTS;
+}
+
+template <typename T>
+struct AggregateFunctionRangeSumData
+{
+    using TimeGap = std::pair<T, T>;
+    using TimeGaps = PODArrayWithStackMemory<TimeGap, 64>;
+
+    bool sorted = false;
+
+    TimeGaps gap_list;
+
+    size_t size() const { return gap_list.size(); }
+
+    void add(T start, T end)
+    {
+        if (sorted && gap_list.size() > 0)
+        {
+            sorted = gap_list.back().first <= start;
+        }
+        gap_list.emplace_back(start, end);
+    }
+
+    void merge(const AggregateFunctionRangeSumData & other)
+    {
+        if (other.gap_list.empty())
+            return;
+
+        const auto size = gap_list.size();
+
+        gap_list.insert(std::begin(other.gap_list), std::end(other.gap_list));
+
+        /// either sort whole container or do so partially merging ranges afterwards
+        if (!sorted && !other.sorted)
+            std::stable_sort(std::begin(gap_list), std::end(gap_list));
+        else
+        {
+            const auto begin = std::begin(gap_list);
+            const auto middle = std::next(begin, size);
+            const auto end = std::end(gap_list);
+
+            if (!sorted)
+                std::stable_sort(begin, middle);
+
+            if (!other.sorted)
+                std::stable_sort(middle, end);
+
+            std::inplace_merge(begin, middle, end);
+        }
+
+        sorted = true;
+    }
+
+    void sort()
+    {
+        if (!sorted)
+        {
+            std::stable_sort(std::begin(gap_list), std::end(gap_list));
+            sorted = true;
+        }
+    }
+
+    void serialize(WriteBuffer & buf) const
+    {
+        writeBinary(sorted, buf);
+        writeBinary(gap_list.size(), buf);
+
+        for (const auto & time_gap : gap_list)
+        {
+            writeBinary(time_gap.first, buf);
+            writeBinary(time_gap.second, buf);
+        }
+    }
+
+    void deserialize(ReadBuffer & buf)
+    {
+        readBinary(sorted, buf);
+
+        size_t size;
+        readBinary(size, buf);
+
+        gap_list.clear();
+        gap_list.reserve(size);
+
+        T start, end;
+
+        for (size_t i = 0; i < size; ++i)
+        {
+            readBinary(start, buf);
+            readBinary(end, buf);
+            gap_list.emplace_back(start, end);
+        }
+    }
+};
+
+template <typename T, typename Data>
+class AggregateFunctionRangeSum final : public IAggregateFunctionDataHelper<Data, AggregateFunctionRangeSum<T, Data>>
+{
+private:
+    UInt64 getRangeSum(Data & data) const
+    {
+        if (data.size() == 0)
+            return 0;
+
+        data.sort();
+
+        typename Data::TimeGaps merged;
+
+        for (const auto & gap : data.gap_list)
+        {
+            if (!merged.size() || merged.back().second < gap.first)
+                merged.emplace_back(gap.first, gap.second);
+            else
+                merged.back().second = std::max(merged.back().second, gap.second);
+        }
+
+        UInt64 res = 0;
+        for (const auto & gap : merged)
+        {
+            res += gap.second - gap.first;
+        }
+
+        return res;
+    }
+
+public:
+    String getName() const override { return "rangeSum"; }
+
+    explicit AggregateFunctionRangeSum(const DataTypes & arguments)
+        : IAggregateFunctionDataHelper<Data, AggregateFunctionRangeSum<T, Data>>(arguments, {})
+    {
+    }
+
+    DataTypePtr getReturnType() const override { return std::make_shared<DataTypeUInt64>(); }
+
+    bool allocatesMemoryInArena() const override { return false; }
+
+    AggregateFunctionPtr getOwnNullAdapter(
+        const AggregateFunctionPtr & nested_function,
+        const DataTypes & arguments,
+        const Array & params,
+        const AggregateFunctionProperties & /*properties*/) const override
+    {
+        return std::make_shared<AggregateFunctionNullVariadic<false, false, false>>(nested_function, arguments, params);
+    }
+
+    void add(AggregateDataPtr __restrict place, const IColumn ** columns, const size_t row_num, Arena *) const override
+    {
+        auto start = assert_cast<const ColumnVector<T> *>(columns[0])->getData()[row_num];
+        auto end = assert_cast<const ColumnVector<T> *>(columns[1])->getData()[row_num];
+        this->data(place).add(start, end);
+    }
+
+    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
+    {
+        this->data(place).merge(this->data(rhs));
+    }
+
+    void serialize(ConstAggregateDataPtr __restrict place, WriteBuffer & buf) const override
+    {
+        this->data(place).serialize(buf);
+    }
+
+    void deserialize(AggregateDataPtr __restrict place, ReadBuffer & buf, Arena *) const override
+    {
+        this->data(place).deserialize(buf);
+    }
+
+    void insertResultInto(AggregateDataPtr __restrict place, IColumn & to, Arena *) const override
+    {
+        assert_cast<ColumnUInt64 &>(to).getData().push_back(getRangeSum(this->data(place)));
+    }
+};
+
+}
diff --git a/src/AggregateFunctions/Helpers.h b/src/AggregateFunctions/Helpers.h
index d1f1830eb27..c2f1434b515 100644
--- a/src/AggregateFunctions/Helpers.h
+++ b/src/AggregateFunctions/Helpers.h
@@ -114,6 +114,17 @@ static IAggregateFunction * createWithUnsignedIntegerType(const IDataType & argu
     return nullptr;
 }
 
+template <template <typename, typename> class AggregateFunctionTemplate, template <typename> class Data, typename... TArgs>
+static IAggregateFunction * createWithNativeUnsignedIntegerType(const IDataType & argument_type, TArgs &&... args)
+{
+    WhichDataType which(argument_type);
+    if (which.idx == TypeIndex::UInt8) return new AggregateFunctionTemplate<UInt8, Data<UInt8>>(std::forward<TArgs>(args)...);
+    if (which.idx == TypeIndex::UInt16) return new AggregateFunctionTemplate<UInt16, Data<UInt16>>(std::forward<TArgs>(args)...);
+    if (which.idx == TypeIndex::UInt32) return new AggregateFunctionTemplate<UInt32, Data<UInt32>>(std::forward<TArgs>(args)...);
+    if (which.idx == TypeIndex::UInt64) return new AggregateFunctionTemplate<UInt64, Data<UInt64>>(std::forward<TArgs>(args)...);
+    return nullptr;
+}
+
 template <template <typename> class AggregateFunctionTemplate, typename... TArgs>
 static IAggregateFunction * createWithNumericBasedType(const IDataType & argument_type, TArgs && ... args)
 {
diff --git a/src/AggregateFunctions/registerAggregateFunctions.cpp b/src/AggregateFunctions/registerAggregateFunctions.cpp
index 7663eb55955..0a94424dbbc 100644
--- a/src/AggregateFunctions/registerAggregateFunctions.cpp
+++ b/src/AggregateFunctions/registerAggregateFunctions.cpp
@@ -62,6 +62,7 @@ void registerAggregateFunctionCombinatorDistinct(AggregateFunctionCombinatorFact
 
 void registerWindowFunctions(AggregateFunctionFactory & factory);
 
+void registerAggregateFunctionRangeSum(AggregateFunctionFactory &);
 
 void registerAggregateFunctions()
 {
@@ -111,6 +112,8 @@ void registerAggregateFunctions()
         registerAggregateFunctionStudentTTest(factory);
 
         registerWindowFunctions(factory);
+
+        registerAggregateFunctionRangeSum(factory);
     }
 
     {