ClickHouse原始碼筆記5:聚合函式的原始碼再梳理

HappenLee 發表於 2021-04-20

筆者在原始碼筆記1之中分析過ClickHouse的聚合函式的實現,但是對於各個介面函式的實際如何共同工作的原始碼,回頭看並沒有那麼明晰,主要原因是沒有結合Aggregator的類來一起分析聚合函式的是如果工作起來的。所以決定重新再完成一篇聚合函式的原始碼梳理的文章,幫助大家進一步的理解ClickHouse之中聚合函式的工作原理。
本系列文章的原始碼分析基於ClickHouse v19.16.2.2的版本。

1.IAggregateFunction介面梳理

話不多說,直接上程式碼,筆者這裡會將所有聚合函式的核心介面程式碼全部列出,一一梳理各個部分:

建構函式
 IAggregateFunction(const DataTypes & argument_types_, const Array & parameters_)
        : argument_types(argument_types_), parameters(parameters_) {}

上面的程式碼實現了IAggregateFunction介面的建構函式,初始化了該介面的兩個成員變數:

  • argument_type:函式的引數型別,比如函式select sum(a), sum(b), c from test group by c, 這裡a, b分別是UInt16型別與Double型別,那麼這個sum(a)sum(b)的引數就不同。
  • parameters: 引數,實際型別為std::vector<Field>。它代表著函式的除了資料的輸入引數之外的其他引數。比如聚合函式topk,其中需要傳入的k的值就在parameters之中。
記憶體分配介面

在Clickhouse的聚合執行過程之中,所有的聚合函式都是通過列來進行的。而這裡有兩個重要的問題:

  • 列記憶體從哪裡分配
  • 分配的記憶體結構,長度是如何的
    筆者在梳理下面程式碼的過程之中給出解答,
    /** Create empty data for aggregation with `placement new` at the specified location.
      * You will have to destroy them using the `destroy` method.
      */
    virtual void create(AggregateDataPtr place) const = 0;

    /// Delete data for aggregation.
    virtual void destroy(AggregateDataPtr place) const noexcept = 0;

IAggregateFunction定義的兩個介面createdestory介面完成了記憶體結構與長度的確定,這裡可能描述的不是很明白,這裡瞭解Doris聚合實現的同學可以這樣理解。create函式本身就是完成了Doris聚合函式之中init函式所完成的工作。這裡通過子類IAggregateFunctionDataHelper的實現程式碼來進一步理解它做了什麼事情:

    void create(AggregateDataPtr place) const override
    {
        new (place) Data;
    }

    void destroy(AggregateDataPtr place) const noexcept override
    {
        data(place).~Data();
    }

這部分程式碼很簡單,Data就是模板派生的型別,然後通過placement newplacement delete的方式完成了Data型別的構造與析構。而這個Data型別就是聚合函式儲存中間結果的型別,比如sum的聚合函式的派生型別是類AggregateFunctionSumData的記憶體結構,它不僅包含了聚合結果的資料sum同時也包含了一組進行聚合計算的函式介面add,merge等:

template <typename T>
struct AggregateFunctionSumData
{
    T sum{};

    void add(T value)
    {
        sum += value;
    }

    void merge(const AggregateFunctionSumData & rhs)
    {
        sum += rhs.sum;
    }

    T get() const
    {
        return sum;
    }
};

這裡就是通過createdestory函式呼叫AggregateFunctionSumData的建構函式與解構函式。而問題又繞回第一個問題了,這部分記憶體是在哪裡分配的呢?

 aggregate_data = aggregates_pool->alignedAlloc(total_size_of_aggregate_states, align_aggregate_states);
 createAggregateStates(aggregate_data);

在進行聚合運算時,通過Aggregator之中的記憶體池進行單行所有的聚合函式的資料結果的記憶體分配。並且呼叫createAggregateStates依次呼叫各個聚合函式的create方法進行建構函式的呼叫。這部分可能有些難理解,我們接著看下面的流程圖,來更好的幫助理解:

create函式在聚合的流程之中的作用

通過上述流程圖可以看到,create這部分就是在構造聚合hash表時,進行記憶體初始化工作的,而這部分記憶體不僅僅包含了聚合函式的結果資料,還包含了對應聚合運算元的函式指標。後文我們分析計算介面的時候也會同樣看到。接下來,來看destory就很容易理解了,就是在聚合計算結束或取消時,遍歷hash表,並呼叫解構函式對hash表中儲存的Data型別呼叫解構函式,而最終的記憶體伴隨著aggregates_pool記憶體池的析構而同時釋放。

detory函式在聚合流程之中的作用

函式計算介面

接下來就是聚合函式最核心的部分,聚合函式的計算。

/** Adds a value into aggregation data on which place points to.
     *  columns points to columns containing arguments of aggregation function.
     *  row_num is number of row which should be added.
     *  Additional parameter arena should be used instead of standard memory allocator if the addition requires memory allocation.
     */
    virtual void add(AggregateDataPtr place, const IColumn ** columns, size_t row_num, Arena * arena) const = 0;

    /// Merges state (on which place points to) with other state of current aggregation function.
    virtual void merge(AggregateDataPtr place, ConstAggregateDataPtr rhs, Arena * arena) const = 0;

    /** Contains a loop with calls to "add" function. You can collect arguments into array "places"
      *  and do a single call to "addBatch" for devirtualization and inlining.
      */
    virtual void addBatch(size_t batch_size, AggregateDataPtr * places, size_t place_offset, const IColumn ** columns, Arena * arena) const = 0;

IAggregateFunction定義的3個介面:

  • add函式將對應AggregateDataPtr指標之中資料取出,與列columns中的第row_num的資料進行對應的聚合計算
  • addBatch函式:這是函式也是非常重要的,雖然它僅僅實現了一個for迴圈呼叫add函式。它通過這樣的方式來減少虛擬函式的呼叫次數,並且增加了編譯器內聯的概率,同樣,它實現了高效的向量化。
  • merge函式:將兩個聚合結果進行合併的函式,通常用在併發執行聚合函式的過程之中,需要將對應的聚合結果進行合併。

這裡的兩個函式類似Doris之中聚合函式的updatemerge。接下來我們看它是如何完成工作的。

首先看聚合節點Aggregetor是如何呼叫addBatch函式:

   /// Add values to the aggregate functions.
    for (AggregateFunctionInstruction * inst = aggregate_instructions; inst->that; ++inst)
        inst->that->addBatch(rows, places.data(), inst->state_offset, inst->arguments, aggregates_pool);

這裡依次遍歷AggregateFunction,並呼叫addBatch介面。而addBatch介面就是一行行的遍歷列,將引數列inst->arguments與上文提到create函式構造的聚合資料結構的兩列列資料進行聚合計算:

    void addBatch(size_t batch_size, AggregateDataPtr * places, size_t place_offset, const IColumn ** columns, Arena * arena) const override
    {
        for (size_t i = 0; i < batch_size; ++i)
            static_cast<const Derived *>(this)->add(places[i] + place_offset, columns, i, arena);
    }

這裡還是呼叫了add函式,我們通過AggregateFunctionSum作為子類來具體看一下add的具體實現:

   void add(AggregateDataPtr place, const IColumn ** columns, size_t row_num, Arena *) const override
    {
        const auto & column = static_cast<const ColVecType &>(*columns[0]);
        this->data(place).add(column.getData()[row_num]);
    }

這裡其實還是呼叫上文提到的AggregateFunctionSumData的記憶體結構的add函式完成聚合計算。而這個add函式就是一個簡單的相加邏輯,這樣就完成了簡單的一次聚合運算。

   void add(T value)
    {
        sum += value;
    }

merge函式的實現邏輯類似於add函式,這裡就不展開再次分析了。

函式結果輸出介面

最後就是聚合函式結果輸出介面,將聚合計算的結果重新組織為列存。

  /// Inserts results into a column.
    virtual void insertResultInto(ConstAggregateDataPtr place, IColumn & to) const = 0;

首先看聚合節點Aggregator是如何呼叫insertResultInto函式的

 data.forEachValue([&](const auto & key, auto & mapped)
    {
        method.insertKeyIntoColumns(key, key_columns, key_sizes);

        for (size_t i = 0; i < params.aggregates_size; ++i)
            aggregate_functions[i]->insertResultInto(
                mapped + offsets_of_aggregate_states[i],
                *final_aggregate_columns[i]);
    });

Aggregetor同樣是遍歷hash表之中的結果,將key列先組織成列存,然後呼叫insertResultInto函式將聚合計算的結果也轉換為列存。
這裡我們找一個sum函式的實現,來看看insertResultInto函式介面是如何工作的:

    void insertResultInto(ConstAggregateDataPtr place, IColumn & to) const override
    {
        auto & column = static_cast<ColVecResult &>(to);
        column.getData().push_back(this->data(place).get());
    }

其實很簡單,就是呼叫AggregateDataPtr,也就是AggregateFunctionSumDataget()函式獲取sum計算的結果,然後新增到列記憶體之中。

get函式介面的實現如下:

    T get() const
    {
        return sum;
    }

2.聚合函式的註冊流程

有了上述的背景知識,我們接下來舉個栗子。來看看一個聚合函式的實現細節,以及它是如何被使用的。

AggregateFunctionSum

這裡選取了一個很簡單的聚合運算元Sum,我們來看看它實現的程式碼細節。
這裡我們可以看到AggregateFunctionSum是個final類,無法被繼承了。而它繼承IAggregateFunctionHelp類與IAggregateFunctionDataHelper類。

  • IAggregateFunctionHelp類 通過CRTP讓父類可以直接呼叫子類的add函式指標而避免了虛擬函式呼叫的開銷。
  • IAggregateFunctionHelper類則包含了Data的模板資料型別,也就是上文提及的AggregateFunctionSumData進行記憶體結構的createdestory等等。

這裡我們就重點看,這個類override了getName方法,返回了對應的名字時sum。並且實現了我們上文提到核心方法。

template <typename T, typename TResult, typename Data>
class AggregateFunctionSum final : public IAggregateFunctionDataHelper<Data, AggregateFunctionSum<T, TResult, Data>>
{
public:
    using ResultDataType = std::conditional_t<IsDecimalNumber<T>, DataTypeDecimal<TResult>, DataTypeNumber<TResult>>;
    using ColVecType = std::conditional_t<IsDecimalNumber<T>, ColumnDecimal<T>, ColumnVector<T>>;
    using ColVecResult = std::conditional_t<IsDecimalNumber<T>, ColumnDecimal<TResult>, ColumnVector<TResult>>;

    String getName() const override { return "sum"; }

    AggregateFunctionSum(const DataTypes & argument_types_)
        : IAggregateFunctionDataHelper<Data, AggregateFunctionSum<T, TResult, Data>>(argument_types_, {})
        , scale(0)
    {}

    AggregateFunctionSum(const IDataType & data_type, const DataTypes & argument_types_)
        : IAggregateFunctionDataHelper<Data, AggregateFunctionSum<T, TResult, Data>>(argument_types_, {})
        , scale(getDecimalScale(data_type))
    {}

    DataTypePtr getReturnType() const override
    {
        if constexpr (IsDecimalNumber<T>)
            return std::make_shared<ResultDataType>(ResultDataType::maxPrecision(), scale);
        else
            return std::make_shared<ResultDataType>();
    }

    void add(AggregateDataPtr place, const IColumn ** columns, size_t row_num, Arena *) const override
    {
        const auto & column = static_cast<const ColVecType &>(*columns[0]);
        this->data(place).add(column.getData()[row_num]);
    }

    void merge(AggregateDataPtr place, ConstAggregateDataPtr rhs, Arena *) const override
    {
        this->data(place).merge(this->data(rhs));
    }

    void serialize(ConstAggregateDataPtr place, WriteBuffer & buf) const override
    {
        this->data(place).write(buf);
    }

    void deserialize(AggregateDataPtr place, ReadBuffer & buf, Arena *) const override
    {
        this->data(place).read(buf);
    }

    void insertResultInto(ConstAggregateDataPtr place, IColumn & to) const override
    {
        auto & column = static_cast<ColVecResult &>(to);
        column.getData().push_back(this->data(place).get());
    }

private:
    UInt32 scale;
};

之前我們講到AggregateFunction的函式就是通過AggregateDataPtr指標來獲取AggregateFunctionSumData的地址,來呼叫add實現聚合運算元的。我們可以看到AggregateFunctionSumData實現了前文提到的add, merge, write,read四大方法,正好與介面IAggregateFunction一一對應上了。

template <typename T>
struct AggregateFunctionSumData
{
    T sum{};

    void add(T value)
    {
        sum += value;
    }

    void merge(const AggregateFunctionSumData & rhs)
    {
        sum += rhs.sum;
    }

    void write(WriteBuffer & buf) const
    {
        writeBinary(sum, buf);
    }

    void read(ReadBuffer & buf)
    {
        readBinary(sum, buf);
    }

    T get() const
    {
        return sum;
    }
};

ClickHouse在Server啟動時。main函式之中會呼叫registerAggregateFunction的初始化函式註冊所有的聚合函式。
然後呼叫到下面的函式註冊sum的聚合函式:

void registerAggregateFunctionSum(AggregateFunctionFactory & factory)
{
    factory.registerFunction("sum", createAggregateFunctionSum<AggregateFunctionSumSimple>, AggregateFunctionFactory::CaseInsensitive);
}

也就是完成了這個sum聚合函式的註冊,後續我們get出來就可以愉快的呼叫啦。(這部分有許多模板派生的複雜程式碼,建議與原始碼結合梳理才能事半功倍~~)

3.要點梳理

第二小節解析了一個聚合函式與介面意義對應的流程,這裡重點梳理聚合函式實現的原始碼要點:

  1. 各個聚合函式核心的實現add,merge與序列化,記憶體結構初始化,記憶體結構釋放的介面。
  2. 各個函式的實現需要繼承IAggregateFunctionDataHelper的介面,而它的父類是IAggregateFunctionHelperIAggregateFunction介面。
  3. ClickHouse的聚合函式保證了每次迴圈遍歷一個Block只呼叫一個IAggregateFunction的聚合函式,這樣最大程度上確保了向量化執行的可能性,減少了資料偏移與依賴。

4. 小結

好了,到這裡也就把ClickHouse聚合函式部分的程式碼梳理完了。
除了sum函式外,其他的函式的執行也是同樣通過類似的方式依次來實現和處理的,原始碼閱讀的步驟也可以參照筆者的分析流程來參考。
筆者是一個ClickHouse的初學者,對ClickHouse有興趣的同學,歡迎多多指教,交流。

5. 參考資料

官方文件
ClickHouse原始碼