Presto 函式開發

McHades發表於2020-07-23
REST函式

0. 寫在前面

Presto Functions 並不能像 Hive UDF 一樣動態載入，需要根據 Function 的型別，實現 Presto 內部定義的不同介面，在 Presto 服務啟動時進行註冊，然後才能在 SQL 執行時進行呼叫。

1. 函式定義

Presto 內部將 Functions 分為以下三大類：

  • Scalar Function，即標量函式。將傳遞給它的一個或者多個引數值，進行計算後，返回一個確定型別的標量值。
  • Aggregation Function，即聚合函式。計算從列中取得的值，返回一個單一的值。
  • Window Function，即開窗函式。計算從分組列中取得的值，並返回多個值。

對於不同型別的函式，需要遵循不同的規則進行實現。

1.1 標量函式

Presto 使用註解框架來實現標量函式，標量函式分別需要定義函式名稱、輸入引數型別和返回結果型別。下面介紹幾種開發標量函式常用的註解：

  • @ScalarFunction：用於宣告標量函式的名稱和別名
  • @Description：用於生成函式的功能描述
  • @SqlType：用於宣告函式的返回型別和引數型別
  • @TypeParameter：用於宣告型別變數，它所宣告的型別變數可以用於函式的返回型別和引數型別，框架在執行時會自動將變數與具體的型別進行繫結
  • @SqlNullable：用於表示函式引數或返回結果可能為NULL。如果方法的引數不使用此註解，當函式引數包含NULL時，則該函式不會被呼叫，框架自動返回結果NULL。當 Java 程式碼中用於實現函式的方法的返回值為包裝型別時，必須要在實現方法上加上該註解，且該註解無法用於 Java 基礎型別

下面用一個簡單的is_null函式來具體說明如何使用以上註解進行標量函式開發。

public class ExampleIsNullFunction
{
    @ScalarFunction(value = "is_null", alias = "isnull")
    @Description("Returns TRUE if the argument is NULL")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNull(@SqlNullable @SqlType(StandardTypes.VARCHAR) Slice string)
    {
        return (string == null);
    }
}

以上程式碼實現的is_null函式功能為：判斷傳入的VARCHAR型別引數是否為NULL，如果為NULL則返回true，否則返回false。其中：

  • @ScalarFunction(value = "is_null", alias = "isnull")宣告瞭函式名為is_null，函式別名為isnull，即在 SQL 中使用is_null和isnull都可以呼叫該函式
  • @Description("Returns TRUE if the argument is NULL")宣告瞭函式描述，使用show functions命令可以看到函式的描述
  • @SqlType(StandardTypes.BOOLEAN)宣告瞭函式的返回型別為BOOLEAN
  • 因為當函式引數為NULL時，我們不能直接返回NULL，而是要進行判斷，所以要加上@SqlNullable避免框架自動返回NULL
  • @SqlType(StandardTypes.VARCHAR)宣告瞭函式的引數型別為VARCHAR

注意到，這裡使用了 Java 型別Slice來接收 SQL 中VARCHAR型別的值。框架會自動將 SQL 中的資料型別與“原生容器型別”（Native container type）進行繫結，目前“原生容器型別”只包括：boolean、long、double、Slice和Block。VARCHAR對應的原生容器型別是Slice而不是String,Slice的本質是對byte[]進行了封裝，為的是更加高效、自由地對記憶體進行操作。Block可以簡單的理解為對應 SQL 中的陣列型別。具體的對應關係和繫結過程涉及 Presto 的型別系統和函式呼叫過程，不是本文講解的重點，故在此不作展開。

進一步地，我們想對 is_null函式進行升級，使它能夠處理任意型別的引數，這時@TypeParameter註解就派上用場了，函式的實現可以改寫為：

@ScalarFunction(value = "is_null", alias = "isnull")
@Description("Returns TRUE if the argument is NULL")
public class ExampleIsNullFunction
{
    private IsNullFunctions()
    {
    }

    @TypeParameter("T")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNullSlice(@SqlNullable @SqlType("T") Slice value)
    {
        return (value == null);
    }

    @TypeParameter("T")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNullLong(@SqlNullable @SqlType("T") Long value)
    {
        return (value == null);
    }

    @TypeParameter("T")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNullDouble(@SqlNullable @SqlType("T") Double value)
    {
        return (value == null);
    }

    @TypeParameter("T")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNullBoolean(@SqlNullable @SqlType("T") Boolean value)
    {
        return (value == null);
    }

    @TypeParameter("T")
    @SqlType(StandardTypes.BOOLEAN)
    public static boolean isNullBlock(@SqlNullable @SqlType("T") Block value)
    {
        return (value == null);
    }
}

可以看到，@TypeParameter的使用有點類似 Java 中泛型的用法，型別變數T在宣告完之後就可以在@SqlType註解中使用。在實際的呼叫過程中，框架會將T與實際 SQL 型別進行繫結，然後再去呼叫以對應的原生容器型別為引數的實際方法。

1.2 聚合函式

聚合的過程一般涉及多行，有一個累積計算的過程，又由於 Presto 是一個分散式的計算引擎，資料分佈在多個節點，所以需要用狀態物件來維護和記錄中間計算結果。

引入狀態之後，Presto 將聚合的過程抽象為三個步驟：

  1. input(state, value)
  2. combine(state1, state2)
  3. output(state, out)

首先，input 階段分別在不同的 worker 中進行，將行值進行累積計算到state中；combine階段將上一步得到的state進行兩兩結合；經過前兩步，最終會得到一個state，在output階段對最終的state進行處理輸出。

在實現方面，聚合函式的開發使用了和標量函式類似的註解框架，但是由於狀態概念的引入，需要定義一個繼承於AccumulatorState介面的狀態介面，對於簡單的聚合，該介面只需要新增聚合所需的getter和setter，框架會自動生成相關的實現和序列化程式碼；如果聚合過程中需要記錄複雜型別（LIST、MAP或自定義的類）的狀態，則需要額外實現AccumulatorStateFactory介面和AccumulatorStateSerializer介面，並在狀態介面上使用@AccumulatorStateMetadata註解，在註解中指定stateFactoryClass和stateSerializerClass。

下面以實現求DOUBLE型別的列均值的聚合函式avg_double為例來說明如何進行簡單聚合函式的開發。

avg_double的聚合狀態只需要記錄累積和與加數個數，所以狀態介面的定義如下：

public interface LongAndDoubleState
        extends AccumulatorState
{
    long getLong();

    void setLong(long value);

    double getDouble();

    void setDouble(double value);
}

使用定義好的狀態介面進行聚合函式實現：

@AggregationFunction("avg_double")
public class AverageAggregation
{
    @InputFunction
    public static void input(LongAndDoubleState state, @SqlType(StandardTypes.DOUBLE) double value)
    {
        state.setLong(state.getLong() + 1);
        state.setDouble(state.getDouble() + value);
    }

    @CombineFunction
    public static void combine(LongAndDoubleState state, LongAndDoubleState otherState)
    {
        state.setLong(state.getLong() + otherState.getLong());
        state.setDouble(state.getDouble() + otherState.getDouble());
    }

    @OutputFunction(StandardTypes.DOUBLE)
    public static void output(LongAndDoubleState state, BlockBuilder out)
    {
        long count = state.getLong();
        if (count == 0) {
            out.appendNull();
        }
        else {
            double value = state.getDouble();
            DOUBLE.writeDouble(out, value / count);
        }
    }
}

可以看到聚合函式的實現使用了以下註解：

  • @AggregationFunction宣告瞭聚合函式的名稱，也可以指定函式的別名
  • @InputFunction、@CombineFunction和@OutputFunction分別用來標記聚合的三個步驟，其中@OutputFunction註解需要宣告聚合函式返回結果的資料型別
  • BlockBuilder類為結果輸出類，聚合計算出的最終結果值將通過BlockBuilder進行輸出

1.3 視窗函式

視窗函式在查詢結果的行上進行計算，執行順序在HAVING子句之後，ORDER BY子句之前。在 Presto SQL 中，視窗函式的語法形式如下：

windowFunction(arg1,....argn) OVER([PARTITION BY<...>] [ORDER BY<...>] [RANGE|ROWS BETWEEN AND])

由此可見，視窗函式語法由關鍵字OVER觸發，且包含三個子句：

  1. PARTITION BY: 指定輸入行分割槽的規則，類似於聚合函式的GROUP BY子句，不同分割槽裡的計算互不干擾(視窗函式的計算是併發進行的，併發數和partition數量一致)，預設時將所有資料行視為一個分割槽
  2. ORDER BY: 決定了視窗函式處理輸入行的順序
  3. RANGE|ROWS BETWEEN AND: 指定視窗邊界，不常用，預設時的視窗為當前行所在的分割槽第一行到當前行

視窗函式的開發需要實現WindowFunction介面，WindowFunction介面中宣告瞭兩個方法：

  • void reset(WindowIndex windowIndex): 處理新分割槽時，都會呼叫該方法進行初始化，WindowIndex包含了已排序的分割槽的所有行
  • void processRow(BlockBuilder output, int peerGroupStart, int peerGroupEnd, int frameStart, int frameEnd): 視窗函式的實現方法，BlockBuilder為結果輸出類，計算出來的值將通過BlockBuilder進行輸出；peerGroupStart和peerGroupEnd為當前處理的行所在的分割槽的開始和結束的位置；frameStart和frameEnd為當前處理行所在的視窗的開始和結束位置。

實現一個返回視窗中第一個值的視窗函式first_value(x)的程式碼如下：

@WindowFunctionSignature(name = "first_value", typeVariable = "T", returnType = "T", argumentTypes = "T")
public class FirstValueFunction
        extends WindowFunction
{
    private final int argumentChannel;
    private WindowIndex windowIndex;

    public FirstValueFunction(List<Integer> argumentChannels)
    {
        this.argumentChannel = getOnlyElement(argumentChannels);
    }

    @Override
    public void reset(WindowIndex windowIndex)
    {
        this.windowIndex = windowIndex;
    }

    @Override
    public void processRow(BlockBuilder output, int peerGroupStart, int peerGroupEnd, int frameStart, int frameEnd)
    {
        if (frameStart < 0) {
            output.appendNull();
            return;
        }

        //Outputs a value from the index
        windowIndex.appendTo(argumentChannel, frameStart, output);
    }
}

其中：

  • @WindowFunctionSignature註解宣告瞭視窗函式的名稱，為了處理任意資料型別的欄位，這裡還宣告瞭型別變數T，並將返回型別和引數型別都指定為T
  • 建構函式中的argumentChannels為引數欄位所在列的索引值
  • processRow方法中，每次只需要通過列索引argumentChannel和當前行所在的視窗起始索引frameStart，就能確定視窗中的第一個值

2. 函式註冊

Presto 函式由MetadataManager中的FunctionRegistry進行管理，開發的函式要生效必須要先註冊到FunctionRegistry中。函式註冊是在 Presto 服務啟動過程中進行的，有以下兩種方式進行函式註冊。

2.1 內建函式註冊

內建函式指的是 Presto 自帶的函式庫中的函式，函式的實現位於presto-main模組中，在FunctionRegistry初始化時進行註冊。具體的註冊過程使用了建造者模式，不同型別的函式註冊只需要呼叫FunctionListBuilder物件對應的方法進行註冊，關鍵程式碼如下：

FunctionListBuilder builder = new FunctionListBuilder()
                .window(RowNumberFunction.class)
                .aggregate(ApproximateCountDistinctAggregation.class)
                .scalar(RepeatFunction.class)
                .function(MAP_HASH_CODE)
                ......

2.2 外掛函式註冊

內建函式滿足不了使用需求時，就需要自行開發函式來擴充函式庫。開發者自行編寫的擴充函式一般通過外掛的方式進行註冊。PluginManager在安裝外掛時會呼叫外掛的getFunctions()方法，將獲取到的函式集合通過MetadataManager的addFunctions方法進行註冊：

public void installPlugin(Plugin plugin)
    {
        ......
       for (Class<?> functionClass : plugin.getFunctions()) {
            log.info("Registering functions from %s", functionClass.getName());
            metadata.addFunctions(extractFunctions(functionClass));
        }
        ......
    }

所以用做擴充函式庫的外掛，需要實現getFunctions()方法，來返回擴充的函式集合，例：

public class ExampleFunctionsPlugin
        implements Plugin
{
    @Override
    public Set<Class<?>> getFunctions()
    {
        return ImmutableSet.<Class<?>>builder()
                .add(ExampleNullFunction.class)
                .add(IsNullFunction.class)
                .add(IsEqualOrNullFunction.class)
                .add(ExampleStringFunction.class)
                .add(ExampleAverageFunction.class)
                .build();
    }
}

3. 多說幾句

以上介紹的 Presto 函式開發方式可以滿足日常大部分函式開發需求， Presto 函式的序號產生器制，新增和修改函式後，必須要重啟服務才能生效，所以目前還不支援真正的使用者自定義函式。

其他較為複雜的函式實現，比如變長引數函式的實現涉及呼叫過程中的函式簽名匹配和型別引數繫結，需要用到codeGen進行實現，具體原理由於篇幅有限，在文中沒有進行展開講解，感興趣的讀者可以在評論區留言。

相關文章