RocksDB 高效采样算法：水塘抽样和随机寻址
SkipListRep::UniqueRandomSample 中实现的这两种非常经典的采样算法。这个函数的设计体现了对算法特性和工程实践的深刻理解，根据不同的场景选择最高效的策略。
它根据采样密度（m 和 N 的关系，其中 m 是 target_sample_size，N 是 num_entries）来选择不同的策略。
场景一：高密度采样 (m > sqrt(N)) - 算法 R (水塘抽样的变种)
当需要采样的数量 m 相对较多时（大于总数 N 的平方根），代码采用了一种非常高效的流式采样算法。
代码分析:
skiplistrep.cc
// ... existing code ...void UniqueRandomSample(const uint64_t num_entries,const uint64_t target_sample_size,std::unordered_set<const char*>* entries) override {
// ... existing code ...// We pick Option 2 when m<sqrt(N) and// Option 1 when m > sqrt(N).if (target_sample_size >static_cast<uint64_t>(std::sqrt(1.0 * num_entries))) {// 对应场景一：高密度采样Random* rnd = Random::GetTLSInstance();iter.SeekToFirst();uint64_t counter = 0, num_samples_left = target_sample_size;for (; iter.Valid() && (num_samples_left > 0); iter.Next(), counter++) {// 核心逻辑：以变化的概率决定是否采样// Add entry to sample set with probability// num_samples_left/(num_entries - counter).if (rnd->Next() % (num_entries - counter) < num_samples_left) {entries->insert(iter.key());num_samples_left--;}}} else {
// ... existing code ...
算法解析:
这个算法被称为算法 R (Algorithm R)，由 Jeffrey Vitter 在其论文《An Efficient Algorithm for Sequential Random Sampling》中提出。它是一种非常巧妙的等概率无放回抽样方法。
- 核心思想: 遍历整个数据集（从
i = 0到N-1），对于第i个元素，以(m - |S|) / (N - i)的概率决定是否将其放入样本集S中。其中|S|是当前样本集的大小。 - 为什么是等概率的? 我们可以通过数学归纳法证明，当这个过程结束后，任何一个元素被选中的概率都是
m/N。- 对于第一个元素，被选中的概率是
m/N。 - 对于第二个元素，它被选中的概率是 (第一个没被选中的概率 * 它被选中的概率) + (第一个被选中的概率 * 它被选中的概率) ... 最终计算结果也是
m/N。
- 对于第一个元素，被选中的概率是
- 优点:
- 只需一次遍历: 算法从头到尾扫描一遍数据即可，时间复杂度为 O(N)。
- 空间高效: 除了存储结果的
m个样本外，不需要额外的巨大存储空间。 - 流式处理: 非常适合处理无法一次性载入内存的大型数据集或数据流。
- 适用场景: 当
m比较大时，随机寻址的方式碰撞概率会很高，导致多次无效寻址。而一次完整的线性遍历反而更划算。m > sqrt(N)这个阈值是一个很好的工程启发式规则。
场景二：低密度采样 (m < sqrt(N)) - 随机寻址法
当需要采样的数量 m 相对较少时，遍历整个数据集 N 就显得非常浪费。此时，代码选择了更直接的方式。
代码分析:
// ... existing code ...} else {// 对应场景二：低密度采样// Option 2: pick m random elements with no duplicates.// If Option 2 is picked, then target_sample_size<sqrt(N)// Using a set spares the need to check for duplicates.for (uint64_t i = 0; i < target_sample_size; i++) {// We give it 5 attempts to find a non-duplicate// ... (概率分析注释) ...for (uint64_t j = 0; j < 5; j++) {iter.RandomSeek();// unordered_set::insert 返回一个 pair，其 .second 成员// 在插入成功时为 true，如果元素已存在则为 false。if ((entries->insert(iter.key())).second) {break;}}}}
// ... existing code ...
算法解析:
- 核心思想: 进行
m次独立的随机采样。每一次都通过iter.RandomSeek()在跳表中随机定位到一个节点，然后尝试将其插入std::unordered_set。 RandomSeek()的实现:InlineSkipList中的RandomSeek是一个 O(logN) 的操作。它从最高层开始，在每一层的节点链表中随机选择一个节点，然后跳到下一层继续这个过程，最终定位到一个随机的元素。unordered_set去重:std::unordered_set的特性被巧妙地利用了。insert()方法在元素不存在时会插入并返回{iterator, true}，如果元素已存在，则不插入并返回{iterator, false}。代码正是通过检查返回的bool值来判断本次随机寻址是否命中了重复元素。- 重试机制: 代码中有一个
for (j < 5)的循环，这是一个为了处理哈希碰撞的重试机制。如果RandomSeek连续几次都找到了已经存在于样本集中的元素，它会重试最多5次。注释中给出了详细的概率分析，证明对于m < sqrt(N)的场景，5次重试足以大概率（>99.9%）找到一个新元素。 - 优点:
- 避免全量扫描: 当
m远小于N时，其期望时间复杂度约为 O(m * logN)，远优于 O(N)。
- 避免全量扫描: 当
- 适用场景: 当采样密度很低时，随机寻址碰到重复元素的概率也很低，这种方法的效率极高。
总结
UniqueRandomSample 函数是自适应算法选择的一个绝佳范例：
| 特性 | 高密度采样 (算法 R) | 低密度采样 (随机寻址) |
|---|---|---|
| 判断条件 | m > sqrt(N) | m <= sqrt(N) |
| 核心操作 | 线性遍历 (iter.Next()) | 随机跳转 (iter.RandomSeek()) |
| 去重方式 | 算法本身保证（通过概率） | std::unordered_set |
| 时间复杂度 | O(N) | 约 O(m * logN) |
| 优势 | m 较大时，避免了随机寻址的高碰撞率，一次遍历解决问题。 | m 较小时，避免了不必要的全量数据扫描，效率极高。 |
通过一个简单的 if-else，RocksDB 就在两种经典的采样算法之间做出了明智的选择，确保了在不同采样需求下都能获得优异的性能。
FindRandomEntry
随机选一层 是一个非常自然的想法，但这种方法无法保证每个元素被选中的概率相等。
高层级的节点本身就比较稀疏，如果随机选一层再随机选一个节点，那么高层级的节点被选中的概率会不成比例地远高于低层级的节点，这就不是一个公平的（均匀的）随机采样了。
FindRandomEntry 的目标是：确保跳表中的每一个元素（即 level 0 上的所有节点）都有完全相同的机会被选中。它通过一个非常巧妙的“分层随机缩放”（Hierarchical Random Zoom-in）策略来实现这一点。
我们来看代码，并用一个具体的例子来解释它的每一步。
inlineskiplist.h
// ... existing code ...
template <class Comparator>
typename InlineSkipList<Comparator>::Node*
InlineSkipList<Comparator>::FindRandomEntry() const {// TODO(bjlemaire): consider adding PREFETCH calls.Node *x = head_, *scan_node = nullptr, *limit_node = nullptr;// ... (注释中的例子) ...std::vector<Node*> lvl_nodes;Random* rnd = Random::GetTLSInstance();int level = GetMaxHeight() - 1;while (level >= 0) {// 1. 收集当前层的候选节点lvl_nodes.clear();scan_node = x;while (scan_node != limit_node) {lvl_nodes.push_back(scan_node);scan_node = scan_node->Next(level);}// 2. 从候选节点中随机选择一个uint32_t rnd_idx = rnd->Next() % lvl_nodes.size();x = lvl_nodes[rnd_idx];// 3. 更新下一轮搜索的范围（"缩放"操作）if (rnd_idx + 1 < lvl_nodes.size()) {limit_node = lvl_nodes[rnd_idx + 1];}level--;}// 4. 处理边界情况并返回结果return x == head_ && head_ != nullptr ? head_->Next(0) : x;
}
// ... existing code ...
分步图解（使用代码注释中的例子）
假设跳表有100个元素，最大高度为5 (GetMaxHeight() 返回 5)。
初始状态:
level = 4(最高层)x = head_(搜索起点是头节点)limit_node = nullptr(搜索终点是整个链表的末尾，没有限制)
第 1 轮: level = 4
- 收集:
while (scan_node != nullptr)循环会遍历第 4 层的所有节点。假设找到了{#1, #15, #67, #84}这几个节点（head_也算一个，但为简化我们从数据节点开始看）。lvl_nodes里面就是这些节点的指针。 - 随机选择:
rnd->Next() % lvl_nodes.size()在这些节点中随机选择一个。假设选中了#15。 - 缩放:
x被更新为#15。limit_node被更新为#15在lvl_nodes中的下一个节点，即#67。- 关键点: 这意味着，下一轮的搜索范围被缩小到了
#15(含) 和#67(不含) 之间！
第 2 轮: level = 3
- 收集:
while (scan_node != #67)循环会遍历第 3 层，但只在#15和#67之间的节点。假设找到了{#15, #21, #45, #51}。 - 随机选择: 在这4个节点中随机选择一个。假设选中了
#51。 - 缩放:
x被更新为#51。limit_node保持为#67，因为#51已经是这个范围内最后一个第3层的节点了。- 搜索范围进一步被缩小到了
#51(含) 和#67(不含) 之间。
第 3 轮: level = 2
- 收集:
while (scan_node != #67)循环遍历第 2 层，范围是#51到#67。假设找到了{#51, #56, #62}。 - 随机选择: 假设选中了
#56。 - 缩放:
x更新为#56。limit_node更新为#62。- 搜索范围缩小为
#56(含) 和#62(不含)。
... 直到 level = 0
最后一轮在 level = 0 上进行，此时的搜索范围可能已经非常小了（比如在 #56 和 #58 之间）。它会在这个极小的范围内再做一次随机选择，最终选出的那个节点就是本次 FindRandomEntry 的结果。
总结
这个算法的精妙之处在于：
- 它不是随机选层，而是逐层随机选择“区间”。
- 在任意一层
L，所有节点把整个链表划分成若干个区间。算法先在L层随机选择一个节点，这等价于随机选择了一个它所代表的区间。 - 然后，算法进入下一层
L-1，但只在刚刚选定的那个小区间内进行下一次的随机选择。 - 这个过程就像从一张世界地图上随机选一个国家，再从这个国家的地图上随机选一个省，再选一个市... 最后定位到一个街道。
通过这种方式，保证了最底层的每一个节点，无论它上面有多少层“快车道”，它被最终选中的概率都是完全一样的 1/N（N为节点总数）。这就实现了一个真正公平、均匀的随机采样。