在數(shù)據(jù)分析領(lǐng)域Python無(wú)疑是最流行的編程語(yǔ)言，但是Python有一個(gè)硬傷就是作為一個(gè)編譯語(yǔ)言在性能上有些微的欠缺。而同樣最流行的語(yǔ)言Rust則在性能方面表現(xiàn)優(yōu)秀。本文我們一起學(xué)習(xí)一個(gè)優(yōu)化項(xiàng)目的實(shí)踐，對(duì)一個(gè)數(shù)據(jù)分析程序，改為Rust后將性能提高了18萬(wàn)倍經(jīng)歷。

概述

要分析的問題如下，以下數(shù)據(jù)是一個(gè)在線問答的數(shù)據(jù)，一個(gè)用戶（user）對(duì)應(yīng)一個(gè)問題（question）以及結(jié)果（score）。

有的用戶可能僅僅回答了問題的一部分，問題的結(jié)果是0或者1。

需要求解問題是：給定一個(gè)大小k, k個(gè)問題的結(jié)合，求解哪一組用戶與整體表現(xiàn)的相關(guān)性最高？

該問題叫做k-CorrSet問題?？梢杂煤?jiǎn)單的簡(jiǎn)單遍歷來解決k-CorrSet問題，算法如下所示（偽代碼）：

func k_corrset($data, $k):$all_qs = all questions in $datafor all $k-sized subsets $qs within $all_qs:$us = all users that answered every question in $qs$qs_totals = the total score on $qs of each user in $us$grand_totals = the grand score on $all_qs of each user in $us$r = correlation($qs_totals, $grand_totals)return $qs with maximum $r

Python算法為基準(zhǔn)

先用Python來解決這個(gè)問題，如果在性能不能滿足需求的話，可以用Rust提高性能。一個(gè)簡(jiǎn)單的Pandas程序來解決k-CorrSet問題的算法：

該算法使用了一些MultiIndex魔法，細(xì)節(jié)上不在深入解釋。馬上進(jìn)行一次開始基準(zhǔn)測(cè)試。

首先，我們需要數(shù)據(jù)。為了使基準(zhǔn)測(cè)試切合實(shí)際，生成了合成數(shù)據(jù)：

60000個(gè)用戶

200個(gè)問題

20%稀疏性（即每個(gè)問題有12,000個(gè)用戶回答）

每個(gè)結(jié)果同樣可能為1或0。

目標(biāo)是計(jì)算該數(shù)據(jù)集上的k-CorrSet，其中k = 5使用2021 M1 Macbook Pro的時(shí)間還算合理。

使用Python的time.time()函數(shù)計(jì)時(shí)，使用 CPython 3.9.17，計(jì)算1000 次迭代的內(nèi)循環(huán)速度。平均執(zhí)行時(shí)間為36毫秒。 還不錯(cuò)，但按照這個(gè)速度，完全完成計(jì)算將在2.9年內(nèi)。

注意：對(duì)Python代碼頁(yè)有很多優(yōu)化技巧，可以提高其性能，如果有需要后續(xù)可以學(xué)習(xí)。

Rust實(shí)現(xiàn)

可以通過將Python代碼用Rust實(shí)現(xiàn)，期待一些免費(fèi)的加速Rust的編譯器優(yōu)化。 為了可讀性，下面的所有代碼都是實(shí)際基準(zhǔn)的簡(jiǎn)化。

首先，轉(zhuǎn)換一下數(shù)據(jù)類型：

在Rust中建立User和Question的新類型，既是為了清晰起見，也是為了在其上使用traits。然后，基本的k-CorrSet算法實(shí)現(xiàn)如下：

fn k_corrset(data: &[Row], k: usize) -> Vec<&Question> {// utils::group_by(impl Iterator<Item = (K1, K2, V)>)// -> HashMap<K1, HashMap<K2, V>>;let q_to_score: HashMap<&Question, HashMap<&User, u32>> =utils::group_by(data.iter().map(|r| (&r.question, &r.user, r.score)));let u_to_score: HashMap<&User, HashMap<&Question, u32>> =utils::group_by(data.iter().map(|r| (&r.user, &r.question, r.score)));let all_grand_totals: HashMap<&User, u32> =u_to_score.iter().map(|(user, scores)| {let total = scores.values().sum::<u32>();(*user, total)}).collect();let all_qs = q_to_score.keys().copied();all_qs.combinations(k).filter_map(|qs: Vec<&Question>| {let (qs_totals, grand_totals): (Vec<_>, Vec<_>) = all_grand_totals.iter().filter_map(|(u, grand_total)| {let q_total = qs.iter().map(|q| q_to_score[*q].get(u).copied()).sum::<Option<u32>>()?;Some((q_total as f64, *grand_total as f64))}).unzip();// utils::correlation(&[f64], &[f64]) -> f64;let r = utils::correlation(&qs_totals, &grand_totals);(!r.is_nan()).then_some((qs, r))}).max_by_key(|(_, r)| FloatOrd(*r)).unwrap().0}

算法關(guān)鍵點(diǎn)：

與Python一樣，將平面數(shù)據(jù)轉(zhuǎn)換為分層數(shù)據(jù)帶有HashMap和utils::group_by幫手。

然后使用Itertools::combinations方法方法迭代所有問題組合。

在內(nèi)循環(huán)中，通過all_grand_totals.iter()方式迭代所有用戶。

表達(dá)方式q_to_score[*q].get(u).copied()有類型 Option<u32>，即 Some(n)如果用戶的結(jié)果為q，否則為None。

如果用戶回答了qs中的所有問題，迭代器方法 .sum::<Option<u32>>()返回Some(total)，否則返回None。

調(diào)用輔助方法utils::correlatio實(shí)現(xiàn)了Pearson的r標(biāo)準(zhǔn)算法。

用max_by_key獲得最高的問題相關(guān)性。用FloatOrd可以比較浮動(dòng)。

那么表現(xiàn)如何呢？ 使用Criterion（默認(rèn)設(shè)置）對(duì)內(nèi)循環(huán)的性能進(jìn)行基準(zhǔn)測(cè)試（filter_map），使用相同的數(shù)據(jù)集。新的內(nèi)循環(huán)運(yùn)行4.2中毫秒，比Python快約8倍基線！

但我們完整的計(jì)算仍然是124天，這有點(diǎn)太長(zhǎng)了。

優(yōu)化

讓我們用一些技巧對(duì)該程序進(jìn)行優(yōu)化一下。

索引數(shù)據(jù)

運(yùn)行一個(gè)探查器，看看程序瓶頸在哪里。 在Mac上，可使用Instruments.app和Samply，后者好像對(duì)Rust優(yōu)化得更好。

下面是用Samply對(duì)Rust算法程序跟蹤相關(guān)部分的屏幕截圖：

可以看到，有75%的時(shí)間都花在HashMap::get上，這是需要優(yōu)化的關(guān)鍵,其對(duì)應(yīng)代碼：

問題是正在散列并比較36字節(jié)UUID字符串，這是一個(gè)昂貴耗時(shí)的操作。對(duì)此，需要一種更小的類型來代替問題/用戶字符串。

解決方案：將所有的問題和用戶收集一個(gè)Vec，并通過索引來表示每個(gè)問題/用戶?？梢允褂胾size指數(shù)與Vec類型，但更好的做法是使用newtypes代表各類指標(biāo)。 事實(shí)上，這個(gè)問題經(jīng)常出現(xiàn)。這樣定義這些索引類型：

QuestionRef和UserRef類型有新類型能夠?qū)崿F(xiàn)traits &Question和&User。define_index_type宏創(chuàng)建新的索引類型QuestionIdx和UserIdx，以及對(duì)應(yīng)的QuestionRef和 UserRef。分別對(duì)應(yīng)為u16和一個(gè)u32類型。

最后更新了k_corrset對(duì)于問題和用戶生成一個(gè)IndexedDomain，然后使用 QuestionIdx和 UserIdx其余代碼中的類型：

我們?cè)俅斡?jì)算的運(yùn)行基準(zhǔn)測(cè)試。新的內(nèi)循環(huán)運(yùn)行時(shí)間為1.0毫秒 ，比上次算法快4，比原始Python版本快35 倍。

總計(jì)算時(shí)間減少到30天，還需要繼續(xù)優(yōu)化。

索引集合

繼續(xù)追蹤執(zhí)行：

仍然，大部分時(shí)間還是消耗在HashMap::get。為了解決這個(gè)問題，考慮完全更換掉HashMap。

HashMap<&User, u32>在概念上和Vec<Option<u32>>是相同的，都對(duì)&User有唯一索引。例如，在一個(gè)Vec中用戶['a', 'b', 'c']，然后是HashMap {'b' => 1}相當(dāng)于vector [None, Some(1), None]。vector消耗更多內(nèi)存，但它改善了鍵/值查找的性能。

考慮到數(shù)據(jù)集規(guī)模進(jìn)行計(jì)算/內(nèi)存權(quán)衡。可以使用Indexical，它提供了 DenseIndexMap<K, V> 內(nèi)部實(shí)現(xiàn)為的類型Vec<V>類型，索引為K::Index。

替換后主要變化是k_corrset函數(shù)，所有輔助數(shù)據(jù)結(jié)構(gòu)轉(zhuǎn)換為DenseIndexMap：

內(nèi)部循環(huán)的唯一變化是：

q_to_score[*q].get(u).copied()

變成了：

再次運(yùn)行基準(zhǔn)測(cè)試，新的內(nèi)循環(huán)運(yùn)行在181微秒 ，比上次迭代快6倍，比原始的Python快了199 倍。

總計(jì)算將縮短至5.3天。

邊界檢查

每次使用括號(hào)時(shí)都會(huì)出現(xiàn)另一個(gè)小的性能影響[]索引到DenseIndexMap。向量 為了安全起見，都要運(yùn)行邊界檢查，實(shí)際上，該代碼可以保證的不會(huì)超出所寫的向量邊界。實(shí)際上找不到邊界檢查樣本配置文件，但它確實(shí)造成了明顯的影響了性能，需要對(duì)其進(jìn)行優(yōu)化。

內(nèi)循環(huán)之前是這樣的：

let q_total = qs.iter().map(|q| q_to_score[*q][u]).sum::<Option<u32>>()?;let grand_total = all_grand_totals[u];

刪除邊界檢查get_unchecked后，新內(nèi)循環(huán)：

沒有邊界檢查是不安全的，所以必須用unsafe塊對(duì)其進(jìn)行標(biāo)記。

再次運(yùn)行基準(zhǔn)測(cè)試，新的內(nèi)循環(huán)運(yùn)行在156微秒，比上一個(gè)迭代快1.16倍，比原始的Python快了229倍。

總計(jì)算將縮短至4.6天。

bit-set

考慮一下內(nèi)循環(huán)的計(jì)算結(jié)構(gòu)。 現(xiàn)在，循環(huán)實(shí)際上看起來像：

for each subset of questions $qs:for each user $u:for each question $q in $qs:if $u answered $q: add $u's score on $q to a running totalelse: skip to the next user$r = correlation($u's totals on $qs, $u's grand total)

數(shù)據(jù)的一個(gè)重要方面是它實(shí)際上形成了一個(gè)稀疏矩陣。對(duì)于給定的問題，只有20%的用戶回答了這個(gè)問題問題。對(duì)于一組5個(gè)問題，只有一小部分回答了全部5個(gè)問題。因此，如果能夠有效地首先確定哪個(gè)用戶回答了所有5個(gè)問題，然后后續(xù)循環(huán)將運(yùn)行減少迭代次數(shù)（并且沒有分支）：

那么我們?nèi)绾伪硎疽鸦卮鸾o定問題的用戶集問題？

可以使用一個(gè)HashSet, 但考慮到到散列的計(jì)算成本很高。因此對(duì)于已索引的數(shù)據(jù)，可以使用更有效的數(shù)據(jù)結(jié)構(gòu)：bit-set，它使用各個(gè)位表示對(duì)象是否存在的內(nèi)存的或集合中不存在。 Indexical提供了另一種抽象將位集與新型索引集成： IndexSet。

此前， q_to_score映射的數(shù)據(jù)結(jié)構(gòu)對(duì)用戶索引的可選分?jǐn)?shù)向量提出問題（即 UserMap<'_, Option<u32>>）。 現(xiàn)在要改變Option<u32>到u32并添加一個(gè)位集描述回答給定問題的一組用戶。 首先更新后的代碼的一半如下所示：

注意q_to_score現(xiàn)在實(shí)際上具有無(wú)效值，因?yàn)闉闆]有回答的用戶提供默認(rèn)值0 問題。

然后更新內(nèi)部循環(huán)以匹配新的偽代碼：

再次運(yùn)行基準(zhǔn)測(cè)試，新的內(nèi)循環(huán)運(yùn)行在47微秒 ，比上次迭代快了3.4倍，比原始Python 程序，快了769倍。

總計(jì)算時(shí)間為1.4天。

單指令多數(shù)據(jù)流

新計(jì)算結(jié)構(gòu)肯定有幫助，但它仍然不夠快。再次檢查一下示例：

現(xiàn)在我們把所有的時(shí)間都花在了bit-set intersection上。因?yàn)槟J(rèn)Indexical使用的位集庫(kù)是bitvec 。其bit-set intersection的原碼是：

fn intersect(dst: &mut BitSet, src: &BitSet) {for (n1, n2): (&mut u64, &u64) in dst.iter_mut().zip(&src) {*n1 &= *n2;}}

bitvec是AND運(yùn)算u64一次。現(xiàn)代大多數(shù)處理器都有專門用于一次執(zhí)行多個(gè)u64位操作指令，稱為SIMD （ingle instruction, multiple data，多數(shù)據(jù)，單指令）。

值得慶幸的是，Rust 提供了實(shí)驗(yàn)性 SIMD API 可以供使用。粗略地說，SIMD版本的bit-set intersection看起來像這樣：

唯一的區(qū)別是已經(jīng)替換了原始的u64類型為SIMD類型u64x4, 在底層，Rust發(fā)出一條SIMD指令來一次執(zhí)行四條u64 &=運(yùn)算。

在crates.io ，搜到一個(gè)名為Bitsvec的?？梢赃m用SIMD的快速交集，但我發(fā)現(xiàn)它的迭代器可以找到索引1位的速度實(shí)際上相當(dāng)慢。進(jìn)行少量修改實(shí)現(xiàn)并編寫了一個(gè)更高效的迭代器。

得益于Indexical的抽象，僅交換SIMD位集需要更改類型別名并且不需要修改k_corrset函數(shù)。優(yōu)化為SIMD位集可以u(píng)64x16在最大程度提高性能。

再次運(yùn)行基準(zhǔn)測(cè)試，新的內(nèi)部循環(huán)運(yùn)行在1.35微秒 ，比上次迭代算法快34倍，比原始Python算法26,459 倍。

總計(jì)算時(shí)間縮短至57分鐘。

內(nèi)存分配

此時(shí)，非常接近峰值性能了。繼續(xù)回到profile倒置視圖（顯示了葉子節(jié)點(diǎn)上最常調(diào)用的函數(shù)調(diào)用樹）：

最大的瓶頸是的位集迭代器。有幾個(gè)相關(guān)的函數(shù)：memmove, realloc,allocate，是在函數(shù)的內(nèi)循環(huán)中分配內(nèi)存的。

為了避免過多分配，可以預(yù)先創(chuàng)建這些數(shù)據(jù)結(jié)構(gòu)所需的最大可能大小，然后重復(fù)寫入他們：

再次運(yùn)行基準(zhǔn)測(cè)試，新的內(nèi)循環(huán)運(yùn)行1.09微秒 ，比上次迭代快1.24倍，比原始的Python基線32,940倍。

總計(jì)算時(shí)間縮短至46分鐘。

并行性

至此，似乎已經(jīng)用盡了所有的優(yōu)化途徑。實(shí)際上想不出任何其他方法來制作內(nèi)循環(huán)速度大大加快。但是實(shí)際上，還可考慮一個(gè)通用技巧并行執(zhí)行！

可以簡(jiǎn)單地并行化內(nèi)部循環(huán)多個(gè)核心運(yùn)行：

let all_qs = questions.indices();all_qs.combinations(k).par_bridge().map_init(|| (vec![0.; users.len()], vec![0.; users.len()], IndexSet::new(&users)),|(qs_totals, grand_totals, user_set), qs| {// same code as before})// same code as before

par_bridge方法采用串行迭代器并且將其轉(zhuǎn)換為并行迭代器。

map_init功能是一個(gè)具有線程特定狀態(tài)的并行映射，所保留免分配狀態(tài)。

需要一個(gè)不同的基準(zhǔn)來評(píng)估外循環(huán)。用5000000個(gè)問題組合上運(yùn)行外循環(huán)的標(biāo)準(zhǔn) 使用給定策略的單次運(yùn)行。

使用串行策略運(yùn)行此基準(zhǔn)測(cè)試超過最快內(nèi)循環(huán)需要6.8秒。對(duì)比并行策略進(jìn)行基準(zhǔn)測(cè)試后，大概需要4.2 秒完成5000000種組合。

只是1.6倍加速

追蹤下性能執(zhí)行：

線程大部分時(shí)間都花在鎖定和解鎖互斥，可能存在某種同步瓶頸。

之間的交接Itertools::combinations迭代器和Rayon并行橋太慢了。鑒于有大量的組合，避免這個(gè)瓶頸的簡(jiǎn)單方法是增加粒度任務(wù)分配。也就是說，可以將許多問題批處理在一起組合并將它們一次性傳遞給一個(gè)線程。

對(duì)于這個(gè)任務(wù)，定義了一個(gè)快速而粗劣的批處理迭代器使用一個(gè)ArrayVec以避免分配。

然后通過批處理組合迭代器來修改外循環(huán)， 并修改內(nèi)部循環(huán)以展平每個(gè)批次：

let all_qs = questions.indices();all_qs.combinations(k).batched::<1024>().par_bridge().map_init(|| (vec![0.; users.len()], vec![0.; users.len()], IndexSet::new(&users)),|(qs_totals, grand_totals, user_set), qs_batch| {qs_batch.into_iter().filter_map(|qs| {// same code as before}).collect_vec()}).flatten()

再次運(yùn)行外循環(huán)基準(zhǔn)測(cè)試，現(xiàn)在是分塊迭代器內(nèi)完成5000000種組合在982毫秒。 與串行方法相比，速度提高了6.9倍。

總結(jié)

結(jié)論

最初的Python程序需要k=5時(shí)需要2.9年完成。使用各種方法優(yōu)化過的Rust程序只需要8 分鐘就可以實(shí)現(xiàn)對(duì)幾十億數(shù)據(jù)的處理?？傮w上，優(yōu)化了180,000 倍加速。

在這個(gè)案例中，使用的優(yōu)化關(guān)鍵點(diǎn)為：

使用Rust的編譯器優(yōu)化。

使用散列數(shù)字而非字符串。

使用（索引）向量而非HashMap。

使用bit-set進(jìn)行有效的成員資格測(cè)試。

使用SIMD實(shí)現(xiàn)高效的位集。

使用多線程將工作分配給多個(gè)核心計(jì)算

使用批處理來避免工作分配中的瓶頸。