1. 函数定义
Trino 内部将 Functions 分为以下三大类:
1、Scalar Function,即标量函数。将传递给它的一个或者多个参数值,进行计算后,返回一个确定类型的标量值。
2、Aggregation Function,即聚合函数。计算从列中取得的值,返回一个单一的值。
3、Window Function,即开窗函数。计算从分组列中取得的值,并返回多个值。
对于不同类型的函数,需要遵循不同的规则进行实现。
1.1 标量函数
我们使用注解框架来实现标量函数,标量函数分别需要定义函数名称、输入参数类型和返回结果类型。下面列出几种开发标量函数常用的注解:
@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。其中:
1、@ScalarFunction(value = "is_null", alias = "isnull")声明了函数名为is_null,函数别名为isnull,即在 SQL 中使用is_null和isnull都可以调用该函数
2、@Description("Returns TRUE if the argument is NULL")声明了函数描述,使用show functions命令可以看到函数的描述
3、@SqlType(StandardTypes.BOOLEAN)声明了函数的返回类型为BOOLEAN
因为当函数参数为NULL时,我们不能直接返回NULL,而是要进行判断,所以要加上@SqlNullable避免框架自动返回NULL
4、@SqlType(StandardTypes.VARCHAR)声明了函数的参数类型为VARCHAR
注意这里使用了 Java 类型Slice来接收 SQL 中VARCHAR类型的值。框架会自动将 SQL 中的数据类型与“原生容器类型”(Native container type)进行绑定,目前“原生容器类型”只包括:boolean、long、double、Slice和Block。VARCHAR对应的原生容器类型是Slice而不是String,Slice的本质是对byte[]进行了封装,为的是更加高效、自由地对内存进行操作。Block可以简单的理解为对应 SQL 中的数组类型。具体的对应关系和绑定过程涉及类型系统和函数调用过程。
进一步地,想对 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 聚合函数
聚合的过程一般涉及多行,有一个累积计算的过程,又由于 Trino 是一个分布式的计算引擎,数据分布在多个节点,所以需要用状态对象来维护和记录中间计算结果。
引入状态之后,Trino 将聚合的过程抽象为三个步骤:
input(state, value)
combine(state1, state2)
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子句之前。窗口函数的语法形式如下:
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接口中声明了两个方法:
1、void reset(WindowIndex windowIndex): 处理新分区时,都会调用该方法进行初始化,WindowIndex包含了已排序的分区的所有行
2、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. 函数注册
函数由MetadataManager中的FunctionRegistry进行管理,函数要生效必须要先注册到FunctionRegistry中。函数注册是在服务启动过程中进行的,有以下两种方式进行函数注册。
2.1 内置函数注册
内置函数指的是 Trino 自带的函数库中的函数,函数的实现位于trino-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();
}
}
上面的流程可以满足大部分函数开发需求, 函数的注册机制,新增和修改函数后,必须要重启服务才能生效,所以目前还不支持真正的用户自定义函数。 其他较为复杂的函数实现,比如变长参数函数的实现涉及调用过程中的函数签名匹配和类型参数绑定,需要用到codeGen进行实现。
2.3 标量函数注册
函数注册实际上是维护FunctinoRegistry类中的一个 MultiMap,Key 为函数的限定名(QualifiedName,可以简单地理解为函数名),Value 为SqlFunction接口的实现类,实际主要为SqlAggregationFunction、SqlWindowFunction和SqlScalarFunction这三个类的子类。SqlScalarFunction是一个抽象类,定义如下:
public abstract class SqlScalarFunction
implements SqlFunction
{
private final Signature signature;
protected SqlScalarFunction(Signature signature)
{
this.signature = requireNonNull(signature, "signature is null");
checkArgument(signature.getKind() == SCALAR, "function kind must be SCALAR");
}
@Override
public final Signature getSignature()
{
return signature;
}
public abstract ScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, TypeManager typeManager, FunctionRegistry functionRegistry);
public static PolymorphicScalarFunctionBuilder builder(Class<?> clazz)
{
return new PolymorphicScalarFunctionBuilder(clazz);
}
}
可以看出,其子类需要获取Signature和实现specialize方法。
首先来看Signature:
public final class Signature
{
private final String name;
private final FunctionKind kind;
private final List<TypeVariableConstraint> typeVariableConstraints;
private final List<LongVariableConstraint> longVariableConstraints;
private final TypeSignature returnType;
private final List<TypeSignature> argumentTypes;
private final boolean variableArity;
....
}
类的成员变量说明如下:
1、name:函数名,不包括参数类型和结果类型,例如:函数isnull(T):boolean的函数名为isnull
2、kind:枚举类型,有 SCALAR、AGGREGATE 和 WINDOW三种取值,用于区分函数类型
3、typeVariableConstraints:类型变量约束,记录函数中的类型变量名,以及类型变量所需要满足的约束条件:类型是否为comparable、orderable 和是否绑定具体类型。例如:contains<T:comparable>(array(T),T):boolean函数要求类型T满足comparable;array_sort<E:orderable>(array(E)):array(E)函数要求类型E满足orderable;判断两个ROW类型是否相等的操作符(操作符也属于标量函数)EQUAL<T:comparable:row<*>>(T,T):boolean要求类型T为ROW类型
4、LongVariableConstraint:长整型变量约束,记录函数中带有约束的长整型变量的计算表达式(一般用于计算返回类型中的长整型变量)。例如:函数concat<u:x + y>(char(x),char(y)):char(u)的返回类型中长整型变量u的计算表达式为x + y
5、returnType:函数的返回类型
6、argumentTypes:函数参数类型
7、variableArity:标记是否为变长参数
以上成员变量都可以从函数实现的类对象中,根据注解规则解析获得。除了获取Signature,由于同一个函数可能会有多个实现(例如上一篇文章介绍的isnull(T):boolean函数,因为传入的参数类型可能不同,所以有五个实现方法),所以还要记录函数的实现方法。源码中将实现方法分为三类:
1、exactimplementation:函数中不包含类型变量,即函数的参数类型和返回类型都是确定的
2、specializedImplementation:函数中包含类型变量,但类型变量作用在具体的 Java 类型(Native Container Type)上
3、genericImplement:函数中包含类型变量,但是类型变量作用在 Object 类型上
Trino 保存的是实现方法的MethodHandle,通过反射获取Method,再保存Method对应的MethodHandle(MethodHandle在JDK1.7引入,调用的效率比反射高),如果该方法不是静态方法,还要将MethodHandle的中的this参数改为Object来避免调用时的类加载问题。所以,抽象方法specialize的本质是通过传入的参数,来获取匹配到的MethodHandle,这部分放到下一节的标量函数调用中进行阐述。
可以看出,标量函数注册的本质是保存函数的Signature和MethodHandle。开发时根据注解框架实现的标量函数,注册时再根据注解解析出Signature和MethodHandle,封装在ParametricScalar对象中。当然,也可以自行继承SqlScalarFunction,自己定义Signature和实现specialize方法。
3. 标量函数调用
标量函数调用的入口为InterpretedFunctionInvoker类的public Object invoke(Signature function, ConnectorSession session, List arguments)方法,形参里的Signature是由语义分析时,根据词法分析得到函数QualifiedName和语法分析得到的参数类型,调用FunctionRegistry中的public Signature resolveFunction(QualifiedName name, List parameterTypes)方法得到。所以,标量函数调用的关键是resolveFunction方法和invoke方法。
首先来看resolveFunction方法,该方法主要通过函数名和函数参数类型来确定Signature,流程如下:
虚线红框中的三个匹配过程实际上是调用了同一个方法:Optional matchFunction(Collection candidates, List parameters, boolean coercionAllowed),其中的coercionAllowed为是否将实参类型转化为形参类型的标识。matchFunction方法等价于为Signature中的变量寻找赋值,不仅要满足变量类型是对应的实际参数类型的超类,而且对应的实际参数还要满足Signature中声明的变量约束。将形参类型和实参进行绑定时,还会做一些约定性的检查:
1、一个类型不能既赋给类型变量(type parameter),又赋给字面变量(literal parameter,如varchar(x)中的x)
2、字面变量不允许跨类型使用
为了便于理解第二个规定,下面例举几个字面变量跨类型使用的例子:
1、x 出现在不同的基本类型中:char(x)和varchar(x)
2、x 出现在同一种基本类型的不同位置:decimal(x,y) 和 decimal(z,x)
3、p 与不同的字面量、类型或者字面变量组合使用:decimal(p,s1) and decimal(p,s2)
还有一个限制是,如果尝试将实际参数类型decimal(1,0)赋给Signature中声明的decimal(x,2),会失败,但是使用decimal(3,1)可以赋值成功。因为根据decimal的定义,precision 必须大于 scale,即x必须大于2。
经过一系列的规则匹配和变量求解,最终会返回一个具体的函数函数签名,签名中的类型都是具体类型(即不含变量)。比如简单 SQLselect isnull('a'),最终得到的Signature是isnull(varchar(1)):boolean,实参中的类型varchar(1)赋给了原先注册的isnull(T):boolean中的类型变量T。
再来看invoke方法,该方法首先会根据传入的Signature调用FunctionRegistry中的getScalarFunctionImplementation来获取最终的MethodHandle,然后使用具体的参数值来进行实际方法的调用(方法中若需要ConnectorSession,也在此进行注入)。因为函数注册维护的是QualifiedName->SqlFunction的映射关系,而调用getScalarFunctionImplementation时传入的Signature并没有记录变量与实参的绑定关系,所以这里需要再进行一次类型变量的求解,这一步的计算其实是可以避免的,因为在resolveFunction中其实已经拿到了变量绑定的关系,可以进行复用,所以340版本中已改为传入带绑定关系的FunctionBinding。
函数注册时说明了一个函数可能有多个实现方法,接下来就是根据形参和实参的绑定关系,调用SqlFunction的specialize方法进行对应参数的 Java 类型的匹配,按照exactimplementation类型->specializedImplementation类型->genericImplement类型的顺序进行匹配,一旦匹配成功则直接返回匹配到的实现方法,如果方法中需要传入依赖变量,也在此步骤中根据绑定关系对MethodHandle进行参数值注入。因为对MethodHandle的反复编译会导致full GC(怀疑是触发了 JVM Bug),所以在FunctionRegistry中为三类函数分别做了个大小为1000,有效时长为1小时的缓存来避免这个问题。
至此,函数的注册和调用的过程已经完成,求解Signature时的类型转换匹配可以作为类型隐式转换的一个入口。