如何用Rust過(guò)程宏魔法簡(jiǎn)化SQL函數(shù)呢？

macro_rules!for_all_cmp_combinations{
($macro:tt$(,$x:tt)*)=>{
$macro!{
[$($x),*],
//comparisonacrossintegertypes
{int16,int32,int32},
{int32,int16,int32},
{int16,int64,int64},
//...

盡管解決了一些問(wèn)題，但這套方案依然有它的痛點(diǎn)：

基于 trait 做類(lèi)型體操使我們不可避免地陷入到與 Rust 編譯器斗智斗勇之中。

依然沒(méi)有全面覆蓋所有可能情況。有相當(dāng)一部分函數(shù)仍然需要開(kāi)發(fā)者手寫(xiě)向量化實(shí)現(xiàn)。

性能。當(dāng)我們需要引入 SIMD 對(duì)部分函數(shù)進(jìn)行優(yōu)化時(shí)，需要重新實(shí)現(xiàn)一套 kernel 函數(shù)。

沒(méi)有對(duì)開(kāi)發(fā)者隱藏全部細(xì)節(jié)。函數(shù)開(kāi)發(fā)者依然需要先熟悉類(lèi)型體操和聲明宏的工作原理，才能比較流暢地添加函數(shù)。

究其原因，我認(rèn)為是函數(shù)的變體形式過(guò)于復(fù)雜，而 Rust 的 trait 和聲明宏系統(tǒng)的靈活性不足導(dǎo)致的。本質(zhì)上是一種元編程能力不夠強(qiáng)大的表現(xiàn)。

2.4 元編程？

讓我們來(lái)看看其他語(yǔ)言和框架是怎么解決這個(gè)問(wèn)題的。

首先是 Python，一種靈活的動(dòng)態(tài)類(lèi)型語(yǔ)言。這是 Flink 中的 Python UDF 接口，其它大數(shù)據(jù)系統(tǒng)的接口也大同小異：

@udf(result_type='BIGINT')
defadd(i,j):
returni+j

我們發(fā)現(xiàn)它是用 @udf 這個(gè)裝飾器標(biāo)記了函數(shù)的簽名信息，然后在運(yùn)行時(shí)對(duì)不同類(lèi)型進(jìn)行相應(yīng)的處理。當(dāng)然，由于它本身是動(dòng)態(tài)類(lèi)型，因此 Rust 中的很多問(wèn)題在 Python 中根本不存在，代價(jià)則是性能損失。

接下來(lái)是 Java，它是一種靜態(tài)類(lèi)型語(yǔ)言，但通過(guò)虛擬機(jī) JIT 運(yùn)行。這是 Flink 中的 Java UDF 接口：

publicstaticclassSubstringFunctionextendsScalarFunction{
publicStringeval(Strings,Integerbegin,Integerend){
returns.substring(begin,end);
}
}

可以看到同樣也很短。這次甚至不需要額外標(biāo)記類(lèi)型了，因?yàn)殪o態(tài)類(lèi)型系統(tǒng)本身就包含了類(lèi)型信息。我們可以通過(guò)運(yùn)行時(shí)反射拿到類(lèi)型信息，并通過(guò) JIT 機(jī)制在運(yùn)行時(shí)生成高效的強(qiáng)類(lèi)型代碼，兼具靈活與性能。

最后是 Zig，一種新時(shí)代的 C 語(yǔ)言。它最大的特色是任何代碼都可以加上 comptime 關(guān)鍵字在編譯時(shí)運(yùn)行，因此具備非常強(qiáng)的元編程能力。tygg 在博文「Zig lang 初體驗(yàn) -- 『大道至簡(jiǎn)』的 comptime[4]」中演示了用 Zig 實(shí)現(xiàn)遲先生類(lèi)型體操的方法：通過(guò) 編譯期反射 和 過(guò)程式的代碼生成 來(lái)代替開(kāi)發(fā)者完成類(lèi)型體操。

用一張表總結(jié)一下：

可以發(fā)現(xiàn)，Zig 語(yǔ)言強(qiáng)大的元編程能力提供了相對(duì)最好的解決方案。

2.5 過(guò)程宏

那么 Rust 里面有沒(méi)有類(lèi)似 Zig 的特性呢。其實(shí)是有的，那就是過(guò)程宏（Procedural Macros）。它可以在編譯期動(dòng)態(tài)執(zhí)行任何 Rust 代碼來(lái)修改 Rust 程序本身。只不過(guò)，它的編譯時(shí)和運(yùn)行時(shí)代碼是物理分開(kāi)的，相比 Zig 的體驗(yàn)沒(méi)有那么統(tǒng)一，但是效果幾乎一樣。

參考 Python UDF 的接口設(shè)計(jì)，我們便得到了 ”大道至簡(jiǎn)“ 的 Rust 函數(shù)接口：

#[function("add(int,int)->int")]
fnadd(a:i32,b:i32)->i32{
a+b
}

從用戶(hù)的角度看，他只需要在自己熟悉的 Rust 函數(shù)上面標(biāo)一個(gè)函數(shù)簽名。其它的類(lèi)型體操和代碼生成操作都被隱藏在過(guò)程宏之后，完全無(wú)需關(guān)心。

此時(shí)我們已經(jīng)拿到了一個(gè)函數(shù)所必須的全部信息，接下來(lái)我們將看到過(guò)程宏如何生成向量化執(zhí)行所需的樣板代碼。

3展開(kāi) #[function]

3.1 解析函數(shù)簽名

首先我們要實(shí)現(xiàn)類(lèi)型反射，也就是分別解析 SQL 函數(shù)和 Rust 函數(shù)的簽名，以此決定后面如何生成代碼。在過(guò)程宏入口處我們會(huì)拿到兩個(gè) TokenStream，分別包含了標(biāo)注信息和函數(shù)本體：

#[proc_macro_attribute]
pubfnfunction(attr:TokenStream,item:TokenStream)->TokenStream{
//attr:"add(int,int)->int"
//item:fnadd(a:i32,b:i32)->i32{a+b}
...
}

我們使用 syn 庫(kù)將 TokenStream 轉(zhuǎn)為 AST，然后：

解析 SQL 函數(shù)簽名字符串，獲取函數(shù)名、輸入輸出類(lèi)型等信息。

解析 Rust 函數(shù)簽名，獲取函數(shù)名、每個(gè)參數(shù)和返回值的類(lèi)型模式、是否 async 等信息。

具體地：

對(duì)于參數(shù)類(lèi)型，我們確定它是 T 或者 Option。

對(duì)于返回值類(lèi)型，我們將其識(shí)別為：T，Option，Result ，Result> 四種類(lèi)型之一。

這將決定我們后面如何調(diào)用函數(shù)以及處理錯(cuò)誤。

3.2 定義類(lèi)型表

作為 trait 類(lèi)型體操的代替方案，我們?cè)谶^(guò)程宏中定義了這樣一張類(lèi)型表，來(lái)描述類(lèi)型系統(tǒng)之間的對(duì)應(yīng)關(guān)系，并且提供了相應(yīng)的查詢(xún)函數(shù)。

//nameprimitivearrayprefixdatatype
constTYPE_MATRIX:&str="
void_NullNull
boolean_BooleanBoolean
smallintyInt16Int16
intyInt32Int32
bigintyInt64Int64
realyFloat32Float32
floatyFloat64Float64
...
varchar_StringUtf8
bytea_BinaryBinary
array_ListList
struct_StructStruct
";

比如當(dāng)我們拿到用戶(hù)的函數(shù)簽名后，

#[function("length(varchar)->int")]

查表即可得知：

第一個(gè)參數(shù) varchar 對(duì)應(yīng)的 array 類(lèi)型為 StringArray

返回值 int 對(duì)應(yīng)的數(shù)據(jù)類(lèi)型為 DataType::Int32，對(duì)應(yīng)的 Builder 類(lèi)型為 Int32Builder

并非所有輸入輸出均為 primitive 類(lèi)型，因此無(wú)法進(jìn)行 SIMD 優(yōu)化

在下面的代碼生成中，這些類(lèi)型將被填入到對(duì)應(yīng)的位置。

3.3 生成求值代碼

在代碼生成階段，我們主要使用 quote 庫(kù)來(lái)生成并組合代碼片段。最終生成的代碼整體結(jié)構(gòu)如下：

quote!{
struct#struct_name;
implScalarFunctionfor#struct_name{
fneval(&self,input:&RecordBatch)->Result{
#downcast_arrays
letmutbuilder=#builder;
#eval
Ok(Arc::new(builder.finish()))
}
}
}

下面我們來(lái)逐個(gè)填寫(xiě)代碼片段，首先是 downcast 輸入 array：

letchildren_indices=(0..self.args.len());
letarrays=children_indices.map(|i|format_ident!("a{i}"));
letarg_arrays=children_indices.map(|i|format_ident!("{}",types::array_type(&self.args[*i])));

letdowncast_arrays=quote!{
#(
let#arrays:&#arg_arrays=input.column(#children_indices).as_any().downcast_ref()
.ok_or_else(||ArrowError::CastError(...))?;
)*
};

builder：

letbuilder_type=format_ident!("{}",types::array_builder_type(ty));
letbuilder=quote!{#builder_type::with_capacity(input.num_rows())};

接下來(lái)是最關(guān)鍵的執(zhí)行部分，我們先寫(xiě)出函數(shù)調(diào)用的那一行：

letinputs=children_indices.map(|i|format_ident!("i{i}"));
letoutput=quote!{#user_fn_name(#(#inputs,)*)};
//example:add(i0,i1)

然后考慮：這個(gè)表達(dá)式返回了什么類(lèi)型呢？這需要根據(jù) Rust 函數(shù)簽名決定，它可能包含 Option，也可能包含 Result。我們進(jìn)行錯(cuò)誤處理，然后將其歸一化到 Option 類(lèi)型：

letoutput=matchuser_fn.return_type_kind{
T=>quote!{Some(#output)},
Option=>quote!{#output},
Result=>quote!{Some(#output?)},
ResultOption=>quote!{#output?},
};
//example:Some(add(i0,i1))

下面考慮：這個(gè)函數(shù)接收什么樣的類(lèi)型作為輸入？這同樣需要根據(jù) Rust 函數(shù)簽名決定，每個(gè)參數(shù)可能是或不是 Option。如果函數(shù)不接受 Option 輸入，但實(shí)際輸入的卻是 null，那么我們默認(rèn)它的返回值就是 null，此時(shí)無(wú)需調(diào)用函數(shù)。因此，我們使用 match 語(yǔ)句來(lái)對(duì)輸入?yún)?shù)做預(yù)處理：

letsome_inputs=inputs.iter()
.zip(user_fn.arg_is_option.iter())
.map(|(input,opt)|{
if*opt{
quote!{#input}
}else{
quote!{Some(#input)}
}
});
letoutput=quote!{
//這里的inputs是從array中拿出來(lái)的Option
match(#(#inputs,)*){
//我們將部分參數(shù)unwrap后再喂給函數(shù)
(#(#some_inputs,)*)=>#output,
//如有unwrap失敗則直接返回null
_=>None,
}
};
//example:
//match(i0,i1){
//(Some(i0),Some(i1))=>Some(add(i0,i1)),
//_=>None,
//}

此時(shí)我們已經(jīng)拿到了一行的返回值，可以將它 append 到 builder 中：

letappend_output=quote!{builder.append_option(#output);};

最后在外面套一層循環(huán)，對(duì)輸入逐行操作：

leteval=quote!{
for(i,(#(#inputs,)*))inmultizip((#(#arrays.iter(),)*)).enumerate(){
#append_output
}
};

如果一切順利的話(huà)，過(guò)程宏展開(kāi)生成的代碼將如 2.1 節(jié)中所示的那樣。

3.4 函數(shù)注冊(cè)

到此為止我們已經(jīng)完成了最核心、最困難的部分，即生成向量化求值代碼。但是，用戶(hù)該怎么使用生成的代碼呢？

注意到一開(kāi)始我們生成了一個(gè) struct。因此，我們可以允許用戶(hù)指定這個(gè) struct 的名稱(chēng)，或者定義一套規(guī)范自動(dòng)生成唯一的名稱(chēng)。這樣用戶(hù)就能在這個(gè) struct 上調(diào)用函數(shù)了。

//指定生成名為Add的struct
#[function("add(int,int)->int",output="Add")]
fnadd(a:i32,b:i32)->i32{
a+b
}

//調(diào)用生成的向量化求值函數(shù)
letinput:RecordBatch=...;
letoutput:RecordBatch=Add.eval(&input).unwrap();

不過(guò)在實(shí)際場(chǎng)景中，很少有這種使用特定函數(shù)的需求。更多是在項(xiàng)目中定義很多函數(shù)，然后在解析 SQL 查詢(xún)時(shí)，動(dòng)態(tài)地查找匹配的函數(shù)。為此我們需要一種全局的函數(shù)注冊(cè)和查找機(jī)制。

問(wèn)題來(lái)了：Rust 本身沒(méi)有反射機(jī)制，如何在運(yùn)行時(shí)獲取所有由 #[function] 靜態(tài)定義的函數(shù)呢？

答案是：利用程序的鏈接時(shí)（link time）特性，將函數(shù)指針等元信息放入特定的 section 中。程序鏈接時(shí)，鏈接器（linker）會(huì)自動(dòng)收集分布在各處的符號(hào)（symbol）集中在一起。程序運(yùn)行時(shí)即可掃描這個(gè) section 獲取全部函數(shù)了。

Rust 社區(qū)的 dtolnay 大佬為此需求做了兩個(gè)開(kāi)箱即用的庫(kù)：linkme[5] 和 inventory[6]。其中前者是直接利用上述機(jī)制，后者是利用 C 標(biāo)準(zhǔn)的 constructor 初始化函數(shù)，但背后的原理沒(méi)有本質(zhì)區(qū)別。下面我們以 linkme 為例來(lái)演示如何實(shí)現(xiàn)注冊(cè)機(jī)制。

首先我們需要在公共庫(kù)（而不是 proc-macro）中定義函數(shù)簽名的結(jié)構(gòu)：

pubstructFunctionSignature{
pubname:String,
pubarg_types:Vec,
pubreturn_type:DataType,
pubfunction:Box,
}

然后定義一個(gè)全局變量 REGISTRY 作為注冊(cè)中心。它會(huì)在第一次被訪(fǎng)問(wèn)時(shí)利用 linkme 將所有 #[function] 定義的函數(shù)收集到一個(gè) HashMap 中：

///Acollectionofdistributed`#[function]`signatures.
#[linkme::distributed_slice]
pubstaticSIGNATURES:[fn()->FunctionSignature];

lazy_static::lazy_static!{
///Globalfunctionregistry.
pubstaticrefREGISTRY:FunctionRegistry={
letmutsignatures=HashMap::>::new();
forsiginSIGNATURES{
letsig=sig();
signatures.entry(sig.name.clone()).or_default().push(sig);
}
FunctionRegistry{signatures}
};
}

最后在 #[function] 過(guò)程宏中，我們?yōu)槊總€(gè)函數(shù)生成如下代碼：

#[linkme::distributed_slice(SIGNATURES)]
fn#sig_name()->FunctionSignature{
FunctionSignature{
name:#name.into(),
arg_types:vec![#(#args),*],
return_type:#ret,
//這里#struct_name就是我們之前生成的函數(shù)結(jié)構(gòu)體
function:Box::new(#struct_name),
}
}

如此一來(lái)，用戶(hù)就可以通過(guò) FunctionRegistry 提供的方法動(dòng)態(tài)查找函數(shù)并進(jìn)行求值了：

letgcd=REGISTRY.get("gcd",&[Int32,Int32],&Int32);
letoutput:RecordBatch=gcd.function.eval(&input).unwrap();

3.5 小結(jié)

以上我們完整闡述了 #[function] 過(guò)程宏的工作原理和實(shí)現(xiàn)過(guò)程：

使用 syn 庫(kù)解析函數(shù)簽名

使用 quote 庫(kù)生成定制化的向量化求值代碼

使用 linkme 庫(kù)實(shí)現(xiàn)函數(shù)的全局注冊(cè)和動(dòng)態(tài)查找

其中：

SQL 簽名決定了如何從 input array 中讀取數(shù)據(jù)，如何生成 output array

Rust 簽名決定了如何調(diào)用用戶(hù)的 Rust 函數(shù)，如何處理空值和錯(cuò)誤

類(lèi)型查找表決定了 SQL 類(lèi)型和 Rust 類(lèi)型的映射關(guān)系

相比 trait + 聲明宏的解決方案，過(guò)程宏中的 “過(guò)程式” 風(fēng)格為我們提供了極大的靈活性，一攬子解決了之前提到的全部問(wèn)題。在下一章中，我們將會(huì)在這個(gè)框架的基礎(chǔ)上繼續(xù)擴(kuò)展，解決更多實(shí)際場(chǎng)景下的復(fù)雜需求。

4高級(jí)功能

抽象的問(wèn)題是簡(jiǎn)單的，但現(xiàn)實(shí)的需求是復(fù)雜的。上面的原型看似解決了所有問(wèn)題，但在 RisingWave 的實(shí)際工程開(kāi)發(fā)中，我們遇到了各種稀奇古怪的需求，都無(wú)法用最原始的 #[function] 宏實(shí)現(xiàn)。下面我們來(lái)逐一介紹這些問(wèn)題，并利用過(guò)程宏的靈活性見(jiàn)招拆招。

4.1 支持多類(lèi)型重載

有些函數(shù)支持大量不同類(lèi)型的重載，例如 + 運(yùn)算對(duì)幾乎支持所有數(shù)字類(lèi)型。此時(shí)我們一般會(huì)復(fù)用同一個(gè)泛型函數(shù)，然后用不同的類(lèi)型去實(shí)例化它。

#[function("add(*int,*int)->auto")]
#[function("add(*float,*float)->auto")]
#[function("add(decimal,decimal)->decimal")]
#[function("add(interval,interval)->interval")]
fnadd(l:T1,r:T2)->Result
where
T1:Into+Debug,
T2:Into+Debug,
T3:CheckedAdd,
{
a.into().checked_add(b.into()).ok_or(ExprError::NumericOutOfRange)
}

因此我們支持在同一個(gè)函數(shù)上同時(shí)標(biāo)記多個(gè)#[function] 宏。此外，我們還支持使用類(lèi)型通配符將一個(gè)#[function] 自動(dòng)展開(kāi)成多個(gè)，并使用 auto 自動(dòng)推斷返回類(lèi)型。例如 *int 通配符表示全部整數(shù)類(lèi)型 int2, int4, int8，那么 add(*int, *int) 將展開(kāi)為 3 x 3 = 9 種整數(shù)的組合，返回值自動(dòng)推斷為兩種類(lèi)型中最大的一個(gè)：

#[function("add(int2,int2)->int2")]
#[function("add(int2,int4)->int4")]
#[function("add(int2,int8)->int8")]
#[function("add(int4,int4)->int4")]
...

而如果泛型不能滿(mǎn)足一些特殊類(lèi)型的要求，你也完全可以定義新函數(shù)進(jìn)行特化（specialization）：

#[function("add(interval,timestamp)->timestamp")]
fninterval_timestamp_add(l:Interval,r:Timestamp)->Result{
r.checked_add(l).ok_or(ExprError::NumericOutOfRange)
}

這一特性幫助我們快速實(shí)現(xiàn)函數(shù)重載，同時(shí)避免了冗余代碼。

4.2 自動(dòng) SIMD 優(yōu)化

作為零開(kāi)銷(xiāo)抽象語(yǔ)言，Rust 從不向性能妥協(xié)，#[function] 宏也是如此。對(duì)于很多簡(jiǎn)單函數(shù)，理論上可以利用 CPU 內(nèi)置的 SIMD 指令實(shí)現(xiàn)上百倍的性能提升。然而，編譯器往往只能對(duì)簡(jiǎn)單的循環(huán)結(jié)構(gòu)實(shí)現(xiàn)自動(dòng) SIMD 向量化。一旦循環(huán)中出現(xiàn)分支跳轉(zhuǎn)等復(fù)雜結(jié)構(gòu)，自動(dòng)向量化就會(huì)失效。

//簡(jiǎn)單循環(huán)支持自動(dòng)向量化
assert_eq!(a.len(),n);
assert_eq!(b.len(),n);
assert_eq!(c.len(),n);
foriin0..n{
c[i]=a[i]+b[i];
}

//一旦出現(xiàn)分支結(jié)構(gòu)，如錯(cuò)誤處理、越界檢查等，自動(dòng)向量化就會(huì)失效
foriin0..n{
c.push(a[i].checked_add(b[i])?);
}

不幸的是，我們前文中生成的代碼結(jié)構(gòu)并不利于編譯器進(jìn)行自動(dòng)向量化，因?yàn)檠h(huán)中的 builder.append_option() 操作本身就自帶條件分支。

為了支持自動(dòng)向量化，我們需要對(duì)代碼生成邏輯進(jìn)一步特化：

首先根據(jù)函數(shù)簽名判斷這個(gè)函數(shù)能否實(shí)現(xiàn) SIMD 優(yōu)化。這需要滿(mǎn)足以下兩個(gè)主要條件：

比如：

#[function("equal(int,int)->boolean")]
fnequal(a:i32,b:i32)->bool{
a==b
}

所有輸入輸出類(lèi)型均為基礎(chǔ)類(lèi)型，即 boolean, int, float, decimal

Rust 函數(shù)的輸入類(lèi)型均不含 Option，輸出不含 Option 和 Result

一旦上述條件滿(mǎn)足，我們會(huì)對(duì) #eval 代碼段進(jìn)行特化，將其替換為這樣的代碼，調(diào)用 arrow-rs 內(nèi)置的 unary 和 binary kernel 實(shí)現(xiàn)自動(dòng)向量化：

//SIMDoptimizationforprimitivetypes
matchself.args.len(){
0=>quote!{
letc=#ret_array_type::from_iter_values(
std::repeat_with(||#user_fn_name()).take(input.num_rows())
);
letarray=Arc::new(c);
},
1=>quote!{
letc:#ret_array_type=arrow_arith::unary(a0,#user_fn_name);
letarray=Arc::new(c);
},
2=>quote!{
letc:#ret_array_type=arrow_arith::binary(a0,a1,#user_fn_name)?;
letarray=Arc::new(c);
},
n=>todo!("SIMDoptimizationfor{n}arguments"),
}

需要注意，如果用戶(hù)函數(shù)本身包含分支結(jié)構(gòu)，那么自動(dòng)向量化也是無(wú)效的。我們只是盡力為編譯器創(chuàng)造了實(shí)現(xiàn)優(yōu)化的條件。另一方面，這一優(yōu)化也不是完全安全的，它會(huì)使得原本為 null 的輸入強(qiáng)制執(zhí)行。例如整數(shù)除法 a / b，如果 b 為 null，原本不會(huì)執(zhí)行，現(xiàn)在卻會(huì)執(zhí)行 a / 0，導(dǎo)致除零異常而崩潰。這種情況下我們只能修改函數(shù)簽名，避免生成特化代碼。

整體而言，實(shí)現(xiàn)這一功能后，用戶(hù)編寫(xiě)代碼不需要有任何變化，但是部分函數(shù)的性能得到了大幅提高。這對(duì)于高性能數(shù)據(jù)處理系統(tǒng)而言是必須的。

4.3 返回字符串直接寫(xiě)入 buffer

很多函數(shù)會(huì)返回字符串。但是樸素地返回 String 會(huì)導(dǎo)致大量動(dòng)態(tài)內(nèi)存分配，降低性能。

#[function("concat(varchar,varchar)->varchar")]
fnconcat(left:&str,right:&str)->String{
format!("{left}{right}")
}

注意到列式內(nèi)存存儲(chǔ)中，StringArray 實(shí)際上是把多個(gè)字符串存放在一段連續(xù)的內(nèi)存上，構(gòu)建這個(gè)數(shù)組的 StringBuilder 實(shí)際上也只是將字符串追加寫(xiě)入同一個(gè) buffer 里。因此函數(shù)返回 String 是沒(méi)有必要的，它可以直接將字符串寫(xiě)入 StringBuilder 的 buffer 中。

于是我們支持對(duì)返回字符串的函數(shù)添加一個(gè) &mut Write 類(lèi)型的 writer 參數(shù)。內(nèi)部可以直接用 write! 方法向 writer 寫(xiě)入返回值。

#[function("concat(varchar,varchar)->varchar")]
fnconcat(left:&str,right:&str,writer:&mutimplstd::Write){
writer.write_str(left).unwrap();
writer.write_str(right).unwrap();
}

在過(guò)程宏的實(shí)現(xiàn)中，我們主要修改了函數(shù)調(diào)用部分：

letwriter=user_fn.write.then(||quote!{&mutbuilder,});
letoutput=quote!{#user_fn_name(#(#inputs,)*#writer)};

以及特化 append_output 的邏輯：

letappend_output=ifuser_fn.write{
quote!{{
if#output.is_some(){//返回值直接在這行寫(xiě)入builder
builder.append_value("");
}else{
builder.append_null();
}
}}
}else{
quote!{builder.append_option(#output);}
};

經(jīng)過(guò)測(cè)試，這一功能也可以大幅提升字符串處理函數(shù)的性能。

4.4 常量預(yù)處理優(yōu)化

有些函數(shù)的某個(gè)參數(shù)往往是一個(gè)常量，并且這個(gè)常量需要經(jīng)過(guò)一個(gè)開(kāi)銷(xiāo)較大的預(yù)處理過(guò)程。這類(lèi)函數(shù)的典型代表是正則表達(dá)式匹配：

//regexp_like(source,pattern)
#[function("regexp_like(varchar,varchar)->boolean")]
fnregexp_like(text:&str,pattern:&str)->Result{
letregex=regex::new(pattern)?;//預(yù)處理：編譯正則表達(dá)式
Ok(regex.is_match(text))
}

對(duì)于一次向量化求值來(lái)說(shuō)，如果輸入的 pattern 是常數(shù)（very likely），那么其實(shí)只需要編譯一次，然后用編譯后的數(shù)據(jù)結(jié)構(gòu)對(duì)每一行文本進(jìn)行匹配即可。但如果不是常數(shù)（unlikely，但是合法行為），則需要對(duì)每一行 pattern 編譯一次再執(zhí)行。

為了支持這一需求，我們修改用戶(hù)接口，將特定參數(shù)的預(yù)處理過(guò)程提取到過(guò)程宏中，然后把預(yù)處理后的類(lèi)型作為參數(shù)：

#[function(
"regexp_like(varchar,varchar)->boolean",
prebuild="Regex::new($1)?"http://$1表示第一個(gè)參數(shù)（下標(biāo)從0開(kāi)始）
)]
fnregexp_like(text:&str,regex:&Regex)->bool{
regex.is_match(text)
}

這樣，過(guò)程宏可以對(duì)這個(gè)函數(shù)生成兩個(gè)版本的代碼：

如果指定參數(shù)為常量，那么在構(gòu)造函數(shù)中執(zhí)行 prebuild 代碼，并將生成的 Regex 中間值存放在 struct 當(dāng)中，在求值階段直接傳入函數(shù)。

如果不是常量，那么在求值階段將 prebuild 代碼嵌入到函數(shù)參數(shù)的位置上。

至于具體的代碼生成邏輯，由于細(xì)節(jié)相當(dāng)復(fù)雜，這里就不再展開(kāi)介紹了。

總之，這一優(yōu)化保證了此類(lèi)函數(shù)各種輸入下都具有最優(yōu)性能，并且極大簡(jiǎn)化了手工實(shí)現(xiàn)的復(fù)雜性。

4.5 表函數(shù)

最后，我們來(lái)看表函數(shù)（Table Function，Postgres 中也稱(chēng) Set-returning Funcion，返回集合的函數(shù)）。這類(lèi)函數(shù)的返回值不再是一行，而是多行。如果同時(shí)返回多列，那么就相當(dāng)于返回一個(gè)表。

select*fromgenerate_series(1,3);
generate_series
-----------------
1
2
3

對(duì)應(yīng)到常見(jiàn)的編程語(yǔ)言中，實(shí)際是一個(gè)生成器函數(shù)（Generator）。以 Python 為例，可以寫(xiě)成這樣：

defgenerate_series(start,end):
foriinrange(start,end+1):
yieldi

Rust 語(yǔ)言目前在 nightly 版本支持生成器，但這一特性尚未 stable。不過(guò)如果不用 yield 語(yǔ)法的話(huà)，我們可以利用 RPIT 特性實(shí)現(xiàn)返回迭代器的函數(shù)，以達(dá)到同樣的效果：

#[function("generate_series(int,int)->setofint")]
fngenerate_series(start:i32,stop:i32)->implIterator{
start..=stop
}

我們支持在 #[function] 簽名中使用 -> setof 以聲明一個(gè)表函數(shù)。它修飾的 Rust 函數(shù)必須返回一個(gè) impl Iterator，其中的 Item 需要匹配返回類(lèi)型。當(dāng)然，Iterator 的內(nèi)外都可以包含 Option 或 Result。

在對(duì)表函數(shù)進(jìn)行向量化求值時(shí)，我們會(huì)對(duì)每一行輸入調(diào)用生成器函數(shù)，然后將每一行返回的多行結(jié)果串聯(lián)起來(lái)，最后按照固定的 chunk size 進(jìn)行切割，依次返回多個(gè) RecordBatch。因此表函數(shù)的向量化接口長(zhǎng)這個(gè)樣子：

pubtraitTableFunction{
fneval(&self,input:&RecordBatch,chunk_size:usize)
->Result>>>;
}

我們給出一組 generate_series 的輸入輸出樣例（假設(shè) chunk size = 2）：

inputoutput
+-------+------++-----+-----------------+
|start|stop||row|generate_series|
+-------+------++-----+-----------------+
|0|0|---->|0|0|
|||+->|2|0|
|0|2|--++-----+-----------------+
+-------+------+|2|1|
|2|2|
+-----+-----------------+

由于表函數(shù)的輸入輸出不再具有一對(duì)一的關(guān)系，我們?cè)?output 中會(huì)額外生成一列row來(lái)表示每一行輸出對(duì)應(yīng) input 中的哪一行輸入。這一關(guān)系信息會(huì)在某些 SQL 查詢(xún)中被使用到。

回到#[function]宏的實(shí)現(xiàn)，它為表函數(shù)生成的代碼實(shí)際上也是一個(gè)生成器。我們?cè)趦?nèi)部使用了futures_async_stream[7]提供的#[try_stream]宏實(shí)現(xiàn) async generator（它依賴(lài) nightly 的 generator 特性），在 stable 版本中則使用genawaiter[8]代替。之所以要使用生成器，則是因?yàn)橐粋€(gè)表函數(shù)可能會(huì)生成非常長(zhǎng)的結(jié)果（例如generate_series(0, 1000000000)），中途必須把控制權(quán)交還調(diào)用者，才能保證系統(tǒng)不被卡死。感興趣的讀者可以思考一下：如果沒(méi)有 generator 機(jī)制，高效的向量化表函數(shù)求值能否實(shí)現(xiàn)？如何實(shí)現(xiàn)？

說(shuō)到這里，多扯兩句。genawaiter 也是個(gè)很有意思的庫(kù)，它使用 async-await 機(jī)制來(lái)在 stable Rust 中實(shí)現(xiàn) generator。我們知道 async-await 本質(zhì)上也是一種 generator，它們都依賴(lài)編譯器的 CPS 變換實(shí)現(xiàn)狀態(tài)機(jī)。不過(guò)出于對(duì)異步編程的強(qiáng)烈需求，async-await 很早就被穩(wěn)定化，而 generator 卻遲遲沒(méi)有穩(wěn)定。由于背后的原理相通，它們可以互相實(shí)現(xiàn)。 此外，目前 Rust 社區(qū)正在積極推動(dòng) async generator 的進(jìn)展，原生的async gen[9]和for await[10]語(yǔ)法剛剛在上個(gè)月進(jìn)入 nightly。不過(guò)由于沒(méi)有和 futures 生態(tài)對(duì)接，整體依然處于不可用狀態(tài)。RisingWave 的流處理引擎就深度依賴(lài) async generator 機(jī)制實(shí)現(xiàn)流算子，以簡(jiǎn)化異步 IO 下的流狀態(tài)管理。不過(guò)這又是一個(gè)龐大的話(huà)題，之后有機(jī)會(huì)再來(lái)介紹這方面的應(yīng)用吧。

5總結(jié)

由于篇幅所限，我們只能展開(kāi)這么多了。如你所見(jiàn)，一個(gè)簡(jiǎn)單的函數(shù)求值背后，隱藏著非常多的設(shè)計(jì)和實(shí)現(xiàn)細(xì)節(jié)：

為了高性能，我們選擇列式內(nèi)存存儲(chǔ)和向量化求值。

存儲(chǔ)數(shù)據(jù)的容器通常是類(lèi)型擦除的結(jié)構(gòu)。但 Rust 是一門(mén)靜態(tài)類(lèi)型語(yǔ)言，用戶(hù)定義的函數(shù)是強(qiáng)類(lèi)型的簽名。這意味著我們需要在編譯期確定每一個(gè)容器的具體類(lèi)型，做類(lèi)型體操來(lái)處理不同類(lèi)型之間的轉(zhuǎn)換，準(zhǔn)確地把數(shù)據(jù)從容器中取出來(lái)喂給函數(shù)，最后高效地將函數(shù)吐出來(lái)的結(jié)果打包回?cái)?shù)據(jù)容器中。

為了將上述過(guò)程隱藏起來(lái)，我們?cè)O(shè)計(jì)了#[function]過(guò)程宏在編譯期做類(lèi)型反射和代碼生成，最終暴露給用戶(hù)一個(gè)盡可能簡(jiǎn)單直觀(guān)的接口。

但是實(shí)際工程中存在各種復(fù)雜需求以及對(duì)性能的要求，我們必須持續(xù)在接口上打洞，并對(duì)代碼生成邏輯進(jìn)行特化。幸好，過(guò)程宏具有非常強(qiáng)的靈活性，使得我們可以敏捷地應(yīng)對(duì)變化的需求。

#[function]宏最初是為 RisingWave 內(nèi)部函數(shù)實(shí)現(xiàn)的一套框架。最近，我們將它從 RisingWave 項(xiàng)目中獨(dú)立出來(lái)，基于 Apache Arrow 標(biāo)準(zhǔn)化成一套通用的用戶(hù)定義函數(shù)接口arrow-udf[11]。如果你的項(xiàng)目也在使用 arrow-rs 進(jìn)行數(shù)據(jù)處理，現(xiàn)在可以直接使用這套#[function]宏定義自己的函數(shù)。如果你在使用 RisingWave，那么從這個(gè)月底發(fā)布的 1.7 版本起，你可以使用這個(gè)庫(kù)來(lái)定義 Rust UDF。它可以編譯成 WebAssembly 模塊插入到 RisingWave 中運(yùn)行。感興趣的讀者也可以閱讀這個(gè)項(xiàng)目的源碼了解更多實(shí)現(xiàn)細(xì)節(jié)。

事實(shí)上，RisingWave 基于 Apache Arrow 構(gòu)建了一整套用戶(hù)定義函數(shù)接口。此前，我們已經(jīng)實(shí)現(xiàn)了服務(wù)器模式的 Python 和 Java UDF。最近，我們又基于 WebAssembly 實(shí)現(xiàn)了 Rust UDF，基于 QuickJS 實(shí)現(xiàn)了 JavaScript UDF。它們都可以嵌入到 RisingWave 中運(yùn)行，以實(shí)現(xiàn)更好的性能和用戶(hù)體驗(yàn)。

審核編輯：劉清