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

概述

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

[
{
"user": "5ea2c2e3-4dc8-4a5a-93ec-18d3d9197374",
"question": "7d42b17d-77ff-4e0a-9a4d-354ddd7bbc57",
"score": 1
},
{
"user": "b7746016-fdbf-4f8a-9f84-05fde7b9c07a",
"question": "7d42b17d-77ff-4e0a-9a4d-354ddd7bbc57",
"score": 0
},
/* ... 跟多數據 ... */
]

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

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

該問題叫做k-CorrSet問題。可以用簡單的簡單遍歷來解決k-CorrSet問題，算法如下所示（偽代碼）：

func k_corrset($data, $k):
$all_qs = all questions in $data
for 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問題的算法：

from itertools import combinations
import pandas as pd
from pandas import IndexSlice as islice
def k_corrset(data, K):
all_qs = data.question.unique()
q_to_score = data.set_index(['question', 'user'])
all_grand_totals = data.groupby('user').score.sum().rename('grand_total')
corrs = []
for qs in combinations(all_qs, K):
qs_data = q_to_score.loc[islice[qs,:],:].swaplevel()
answered_all = qs_data.groupby(level=[0]).size() == K
answered_all = answered_all[answered_all].index
qs_totals = qs_data.loc[islice[answered_all,:]] 
.groupby(level=[0]).sum().rename(columns={'score': 'qs'})
r = qs_totals.join(all_grand_totals).corr().qs.grand_total
corrs.append({'qs': qs, 'r': r})
corrs = pd.DataFrame(corrs)
return corrs.sort_values('r', ascending=False).iloc[0].qs
data = pd.read_json('scores.json')
print(k_corrset(data, K=5))

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

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

60000個用戶

200個問題

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

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

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

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

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

Rust實現

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

首先，轉換一下數據類型：

pub struct User(pub String);
pub struct Question(pub String);
pub struct Row {
pub user: User,
pub question: Question,
pub score: u32,
}

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

fn k_corrset(data: &[Row], k: usize) -> Vec<&Question> {
// utils::group_by(impl Iterator)
// -> HashMap>;
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::();
(*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::>()?;
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一樣，將平面數據轉換為分層數據帶有HashMap和utils::group_by幫手。

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

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

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

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

調用輔助方法utils::correlatio實現了Pearson的r標準算法。

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

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

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

逐步優化

讓我們用一些技巧對該程序進行優化一下。

索引數據

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

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

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

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

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

解決方案：將所有的問題和用戶收集一個Vec，并通過索引來表示每個問題/用戶。可以使用usize指數與Vec類型，但更好的做法是使用newtypes代表各類指標。事實上，這個問題經常出現。這樣定義這些索引類型：

pub struct QuestionRef<'a>(pub &'a Question);
pub struct UserRef<'a>(pub &'a User);
define_index_type! {
pub struct QuestionIdx for QuestionRef<'a> = u16;
}
define_index_type! {
pub struct UserIdx for UserRef<'a> = u32;
}

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

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

fn k_corrset(data: &[Row], k: usize) -> Vec<&Question> {
let (questions_set, users_set): (HashSet<_>, HashSet<_>) = data.iter()
.map(|row| (QuestionRef(&row.question), UserRef(&row.user)))
.unzip();
let questions = IndexedDomain::from_iter(questions_set);
let users = IndexedDomain::from_iter(users_set);
let q_to_score: HashMap> =
utils::group_by(data.iter().map(|r| (
questions.index(&(QuestionRef(&r.question))),
users.index(&(UserRef(&r.user))),
r.score,
)));
let u_to_score: HashMap> =
utils::group_by(data.iter().map(|r| (
users.index(&(UserRef(&r.user))),
questions.index(&(QuestionRef(&r.question))),
r.score,
)));
let all_grand_totals = // same code
let all_qs = questions.indices();
all_qs.combinations(k)
.filter_map(|qs: Vec| {
})
.max_by_key(|(_, r)| FloatOrd(*r))
.unwrap().0
.into_iter().map(|idx| questions.value(idx).0).collect()
}

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

總計算時間減少到30天，還需要繼續優化。

索引集合

繼續追蹤執行：

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

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

考慮到數據集規模進行計算/內存權衡。可以使用Indexical，它提供了 DenseIndexMap 內部實現為的類型Vec類型，索引為K::Index。

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

pub type QuestionMap<'a, T> = DenseIndexMap<'a, QuestionRef<'a>, T>;
pub type UserMap<'a, T> = DenseIndexMap<'a, UserRef<'a>, T>;
fn k_corrset(data: &[Row], k: usize) -> Vec<&Question> {
let mut q_to_score: QuestionMap<'_, UserMap<'_, Option>> =
QuestionMap::new(&questions, |_| UserMap::new(&users, |_| None));
for r in data {
q_to_score
.get_mut(&QuestionRef(&r.question))
.unwrap()
.insert(UserRef(&r.user), Some(r.score));
}
let grand_totals = UserMap::new(&users, |u| {
q_to_score.values().filter_map(|v| v[u]).sum::()
});
let all_qs = questions.indices();
all_qs.combinations(k)
}

內部循環的唯一變化是：

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

變成了：

q_to_score[*q][u]

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

總計算將縮短至5.3天。

邊界檢查

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

內循環之前是這樣的：

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

刪除邊界檢查get_unchecked后，新內循環：

let q_total = qs.iter()
.map(|q| unsafe {
let u_scores = q_to_score.get_unchecked(q);
*u_scores.get_unchecked(u)
})
.sum::>()?;
let grand_total = unsafe { *all_grand_totals.get_unchecked(u) };

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

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

總計算將縮短至4.6天。

bit-set

考慮一下內循環的計算結構。現在，循環實際上看起來像：

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 total
else: skip to the next user
$r = correlation($u's totals on $qs, $u's grand total)

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

for each subset of questions $qs:
$qs_u = all users who have answered every question in $qs
for each user $u in $qs_u:
for each question $q in $qs:
add $u's score on $q to a running total
$r = correlation($u's scores on $qs, $u's grand total)

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

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

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

type UserSet<'a> = IndexSet<'a, UserRef<'a>>;
let mut q_to_score: QuestionMap<'_, (UserSet<'_>, UserMap<'_, u32>)> =
QuestionMap::new(&questions, |_| (
UserMap::<'_, u32>::new(&users, |_| 0),
UserSet::new(&users),
));
for r in data {
let (scores, set) = &mut q_to_score.get_mut(&QuestionRef(&r.question)).unwrap();
scores.insert(UserRef(&r.user), r.score);
set.insert(UserRef(&r.user));
}

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

然后更新內部循環以匹配新的偽代碼：

let all_qs = questions.indices();
all_qs.combinations(k)
.filter_map(|qs: Vec| {
// Compute the intersection of the user-sets for each question
let mut users = q_to_score[qs[0]].1.clone();
for q in &qs[1..] {
users.intersect(&q_to_score[*q].1);
}
let (qs_totals, grand_totals): (Vec<_>, Vec<_>) = users.indices()
// only .map, not .filter_map as before
.map(|u| {
let q_total = qs.iter()
.map(|q| unsafe {
let (u_scores, _) = q_to_score.get_unchecked(q);
*u_scores.get_unchecked(u)
})
// only u32, not Option as before
.sum::();
let grand_total = unsafe { *all_grand_totals.get_unchecked(u) };
(q_total as f64, grand_total as f64)
})
.unzip();
let r = utils::correlation(&qs_totals, &grand_totals);
(!r.is_nan()).then_some((qs, r))
})

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

總計算時間為1.4天。

單指令多數據流

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

現在我們把所有的時間都花在了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一次。現代大多數處理器都有專門用于一次執行多個u64位操作指令，稱為SIMD （ingle instruction, multiple data，多數據，單指令）。

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

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

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

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

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

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

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

內存分配

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

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

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

let mut qs_totals = vec![0.; users.len()]
let mut grand_totals = vec![0.; users.len()];
let mut user_set = IndexSet::new(&users);
let all_qs = questions.indices();
all_qs.combinations(k)
.filter_map(|qs| {
// Use `clone_from` rather than `clone` to copy without allocation
user_set.clone_from(&q_to_score[qs[0]].1);
for q in &qs[1..] {
user_set.intersect(&q_to_score[*q].1);
}
let mut n = 0;
for (i, u) in user_set.indices().enumerate() {
let q_total = qs.iter()
.map(|q| unsafe {
let (u_scores, _) = q_to_score.get_unchecked(q);
*u_scores..get_unchecked(u)
})
.sum::();
let grand_total = unsafe { *all_grand_totals.get_unchecked(u) };
unsafe {
*qs_totals.get_unchecked_mut(i) = q_total as f64;
*grand_totals.get_unchecked_mut(i) = grand_total as f64;
}
n += 1;
}
let r = utils::correlation(&qs_totals[..n], &grand_totals[..n]);
(!r.is_nan()).then_some((qs, r))
})

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

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

并行性

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

可以簡單地并行化內部循環多個核心運行：

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功能是一個具有線程特定狀態的并行映射，所保留免分配狀態。

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

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

只是1.6倍加速

追蹤下性能執行：

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

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

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

pub struct Batched {
iter: I,
}
impl Iterator for Batched {
type Item = ArrayVec;
#[inline]
fn next(&mut self) -> Option {
let batch = ArrayVec::from_iter((&mut self.iter).take(N));
(!batch.is_empty()).then_some(batch)
}
}

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

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

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

總結

結論

最初的Python程序需要k=5時需要2.9年完成。使用各種方法優化過的Rust程序只需要8 分鐘就可以實現對幾十億數據的處理。總體上，優化了180,000 倍加速。

在這個案例中，使用的優化關鍵點為：

使用Rust的編譯器優化。

使用散列數字而非字符串。

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

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

使用SIMD實現高效的位集。

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

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

審核編輯：湯梓紅