在數(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)歷。

概述

要分析的問(wèn)題如下，以下數(shù)據(jù)是一個(gè)在線問(wèn)答的數(shù)據(jù)，一個(gè)用戶（user）對(duì)應(yīng)一個(gè)問(wèn)題（question）以及結(jié)果（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
},
/* ... 跟多數(shù)據(jù) ... */
]

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

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

該問(wèn)題叫做k-CorrSet問(wèn)題。可以用簡(jiǎn)單的簡(jiǎn)單遍歷來(lái)解決k-CorrSet問(wèn)題，算法如下所示（偽代碼）：

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算法為基準(zhǔn)

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

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魔法，細(xì)節(jié)上不在深入解釋。馬上進(jìn)行一次開(kāi)始基準(zhǔn)測(cè)試。

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

60000個(gè)用戶

200個(gè)問(wèn)題

20%稀疏性（即每個(gè)問(wèn)題有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)

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

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

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的新類型，既是為了清晰起見(jiàn)，也是為了在其上使用traits。然后，基本的k-CorrSet算法實(shí)現(xiàn)如下：

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
}

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

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

然后使用Itertools::combinations方法方法迭代所有問(wèn)題組合。

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

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

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

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

用max_by_key獲得最高的問(wèn)題相關(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)代碼：

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

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

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

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類型有新類型能夠?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ì)于問(wèn)題和用戶生成一個(gè)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()
}

我們?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è)問(wèn)題，考慮完全更換掉HashMap。

HashMap<&User, u32>在概念上和Vec>是相同的，都對(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)衡?？梢允褂肐ndexical，它提供了 DenseIndexMap 內(nèi)部實(shí)現(xiàn)為的類型Vec類型，索引為K::Index。

替換后主要變化是k_corrset函數(shù)，所有輔助數(shù)據(jù)結(jié)構(gòu)轉(zhuǎn)換為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)
}

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

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

變成了：

q_to_score[*q][u]

再次運(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。向量 為了安全起見(jiàn)，都要運(yùn)行邊界檢查，實(shí)際上，該代碼可以保證的不會(huì)超出所寫(xiě)的向量邊界。實(shí)際上找不到邊界檢查樣本配置文件，但它確實(shí)造成了明顯的影響了性能，需要對(duì)其進(jìn)行優(yōu)化。

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

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

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

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) };

沒(méi)有邊界檢查是不安全的，所以必須用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í)際上看起來(lái)像：

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)

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

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)

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

可以使用一個(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ù)向量提出問(wèn)題（即 UserMap<'_, Option>）?，F(xiàn)在要改變Option到u32并添加一個(gè)位集描述回答給定問(wèn)題的一組用戶。首先更新后的代碼的一半如下所示：

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現(xiàn)在實(shí)際上具有無(wú)效值，因?yàn)闉闆](méi)有回答的用戶提供默認(rèn)值0 問(wèn)題。

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

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

再次運(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ù)處理器都有專門(mén)用于一次執(zhí)行多個(gè)u64位操作指令，稱為SIMD （ingle instruction, multiple data，多數(shù)據(jù)，單指令）。

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

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

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

在crates.io ，搜到一個(gè)名為Bitsvec的。可以適用SIMD的快速交集，但我發(fā)現(xiàn)它的迭代器可以找到索引1位的速度實(shí)際上相當(dāng)慢。進(jìn)行少量修改實(shí)現(xiàn)并編寫(xiě)了一個(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)用樹(shù)）：

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

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

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

再次運(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í)際上想不出任何其他方法來(lái)制作內(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)來(lái)評(píng)估外循環(huán)。用5000000個(gè)問(wèn)題組合上運(yùn)行外循環(huán)的標(biāo)準(zhǔn) 使用給定策略的單次運(yùn)行。

使用串行策略運(yùn)行此基準(zhǔn)測(cè)試超過(guò)最快內(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ù)分配。也就是說(shuō)，可以將許多問(wèn)題批處理在一起組合并將它們一次性傳遞給一個(gè)線程。

對(duì)于這個(gè)任務(wù)，定義了一個(gè)快速而粗劣的批處理迭代器使用一個(gè)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)
}
}

然后通過(guò)批處理組合迭代器來(lái)修改外循環(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)化過(guò)的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ì)算

使用批處理來(lái)避免工作分配中的瓶頸。

審核編輯：湯梓紅