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

概述

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

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

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

該問題叫做k-CorrSet問題?？梢杂煤唵蔚暮唵伪闅v來解決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算法為基準

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

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

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

60000個用戶

200個問題

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

每個結果同樣可能為1或0。

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

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

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

Rust實現(xiàn)

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

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

在Rust中建立User和Question的新類型，既是為了清晰起見，也是為了在其上使用traits。然后，基本的k-CorrSet算法實現(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}

算法關鍵點：

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

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

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

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

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

調用輔助方法utils::correlatio實現(xiàn)了Pearson的r標準算法。

用max_by_key獲得最高的問題相關性。用FloatOrd可以比較浮動。

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

但我們完整的計算仍然是124天，這有點太長了。

優(yōu)化

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

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

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

下面是用Samply對Rust算法程序跟蹤相關部分的屏幕截圖：

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

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

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

QuestionRef和UserRef類型有新類型能夠實現(xiàn)traits &Question和&User。define_index_type宏創(chuàng)建新的索引類型QuestionIdx和UserIdx，以及對應的QuestionRef和 UserRef。分別對應為u16和一個u32類型。

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

我們再次計算的運行基準測試。新的內循環(huán)運行時間為1.0毫秒 ，比上次算法快4，比原始Python版本快35 倍。

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

索引集合

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

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

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

考慮到數(shù)據(jù)集規(guī)模進行計算/內存權衡?？梢允褂肐ndexical，它提供了 DenseIndexMap<K, V> 內部實現(xiàn)為的類型Vec<V>類型，索引為K::Index。

替換后主要變化是k_corrset函數(shù)，所有輔助數(shù)據(jù)結構轉換為DenseIndexMap：

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

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

變成了：

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

總計算將縮短至5.3天。

邊界檢查

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

內循環(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后，新內循環(huán)：

沒有邊界檢查是不安全的，所以必須用unsafe塊對其進行標記。

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

總計算將縮短至4.6天。

bit-set

考慮一下內循環(huán)的計算結構。 現(xiàn)在，循環(huán)實際上看起來像：

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

那么我們如何表示已回答給定問題的用戶集問題？

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

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

注意q_to_score現(xiàn)在實際上具有無效值，因為為沒有回答的用戶提供默認值0 問題。

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

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

總計算時間為1.4天。

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

新計算結構肯定有幫助，但它仍然不夠快。再次檢查一下示例：

現(xiàn)在我們把所有的時間都花在了bit-set intersection上。因為默認Indexical使用的位集庫是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運算u64一次?，F(xiàn)代大多數(shù)處理器都有專門用于一次執(zhí)行多個u64位操作指令，稱為SIMD （ingle instruction, multiple data，多數(shù)據(jù)，單指令）。

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

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

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

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

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

總計算時間縮短至57分鐘。

內存分配

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

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

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

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

總計算時間縮短至46分鐘。

并行性

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

可以簡單地并行化內部循環(huá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方法采用串行迭代器并且將其轉換為并行迭代器。

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

需要一個不同的基準來評估外循環(huán)。用5000000個問題組合上運行外循環(huán)的標準 使用給定策略的單次運行。

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

只是1.6倍加速

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

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

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

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

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

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()

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

總結

結論

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

在這個案例中，使用的優(yōu)化關鍵點為：

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

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

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

使用bit-set進行有效的成員資格測試。

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

使用多線程將工作分配給多個核心計算

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