使用SQL进行采样

采样是从大型数据集中提取意义的最强大工具之一。它可以让你将海量数据缩减为一个小而具有代表性的数据集,使用起来快速而简单。

如果你知道如何使用SQL(这种无处不在的查询语言)进行采样,你就可以在任何地方进行采样。没有任何数据集会超出你的掌控!

在这篇文章中,我们将看一些巧妙的采样算法。这些算法速度快,而且很容易转换成SQL。

不过,首先我要指出,许多数据库系统都内置了一些采样支持。例如,一些SQL方言支持TABLESAMPLE子句。如果你的系统有内置支持 - 而且它能满足你的需求 - 使用它通常是你的最佳选择。

然而,内置支持通常仅限于简单情况。让我们考虑一些更具挑战性的现实场景:

• 我们希望能够进行有放回和无放回采样。
• 我们希望进行加权采样,其中输入数据集中的每个项目被选中的概率与其相应的权重成正比。
• 我们希望支持在FAANG规模的数据集中可能看到的全范围权重,比如频率分布(即点击、展示或RPC事件)的10^0到10^9,以及归一化概率分布的10^-9到10^0,值小至10^-324。换句话说,权重是非负数,可能非常大或非常小。
• 我们希望进行确定性采样。这个特性让我们可以进行可重复的采样,在某些情况下,还可以帮助查询规划器生成更快的查询。

无放回采样:SQL中的A-ES算法

2006年,Pavlos S. Efraimidis和Paul G. Spirakis发表了一个一次遍历算法,用于从带权重项目的总体中进行无放回的加权随机采样。它非常简单:

给定一个由i索引的总体,权重为w_i:

1. 对每个i,令k_i = u_i^(1/w_i),其中u_i是(0,1)上的均匀随机数。
2. 选择k_i值最大的n个项目。

该算法在SQL中有一个直接的实现:

SELECT *
FROM Population
ORDER BY POWER(RANDOM(), 1.0 / weight) DESC
LIMIT n

你会注意到我们稍微改变了排序逻辑。直接翻译应该是:

SELECT *
FROM Population
ORDER BY POWER(RANDOM(), 1.0 / weight) DESC
LIMIT n

我们的翻译:

SELECT *
FROM Population
ORDER BY -LN(RANDOM()) / weight
LIMIT n

在数值上更稳定,也有助于展示该算法与迷人的泊松过程世界之间的联系。这种联系使得更容易理解该算法的工作原理。稍后我们会详细讨论这一点。

数值稳定性和其他调整

首先,关于数值稳定性的说法。假设我们的SQL系统在底层使用IEEE双精度浮点值。当权重很大时,比如在10^9数量级,随机值是多少并不重要。相应的排序键几乎总是10^-9。考虑区间[0.01, 0.99],代表99%可能的随机值u。当w = 10^9时,这整个区间都映射到10^-9:

u^(1/w) = 0.01^(1/1e9) ≈ 0.9999999999
u^(1/w) = 0.99^(1/1e9) ≈ 0.9999999999

同样,当权重很小时,比如10^-9,相应的排序键几乎总是零。考虑区间[0.01, 0.99],代表99%可能的随机值u:

u^(1/w) = 0.01^(1/1e-9) ≈ 0
u^(1/w) = 0.99^(1/1e-9) ≈ 0

对于非常大(或小)的权重,直接实现就不起作用了。当非常大(或小)的幂导致本应不同的排序键崩塌成无法区分的固定点时,所需的排序顺序就被破坏了。

幸运的是,对数是保序变换,所以当我们使用数学上纯粹的实数时,按-ln(u)/w排序产生的顺序与按u^(1/w)排序相同。但是在使用浮点数时,对数变换版本要稳定得多。即使输入权重非常大或非常小,不同的随机输入现在也能可靠地产生不同的排序键-ln(u)/w:

-ln(0.01) / 1e9 ≈ 4.6e-9
-ln(0.99) / 1e9 ≈ 1.0e-11
-ln(0.01) / 1e-9 ≈ 4.6e9
-ln(0.99) / 1e-9 ≈ 1.0e7

最后一个调整是,我们对排序键取反,这样我们就可以按升序排序,而不是像原算法那样按降序排序。注意前面的负号。翻转符号的理由将在下一节讨论泊松过程时变得明显。

最后一个数值上的微妙之处。为什么我们用表达式1.0 - RANDOM()而不是直接用RANDOM()来生成随机数?因为大多数RANDOM()的实现,比如DuckDB使用的PCG实现,返回半开区间[0,1)上的浮点值,理论上可能返回零。而我们不想除以零。所以我们用1.0 - RANDOM()来生成半开区间(0,1]上的随机数,这就排除了零。

这个算法真的有效吗?

给每一行分配一个随机数r_i,然后按-ln(r_i)/w_i对行进行排序,最后取前n行,这个方法能得到有效样本并不明显。但它确实有效。

理解其中原理最清晰的方法是先绕道看看迷人的泊松过程世界。简而言之,具有率λ的泊松过程是一个到达序列,其中连续到达之间的时间都是独立的、具有率λ的指数分布随机变量。

泊松过程有一些重要(且有用!)的性质:

1. 它们是无记忆的。 无论之前发生了什么,从具有率λ的泊松过程到下一次到达的时间都是具有率λ的指数分布。
2. 它们可以合并。 如果你有两个泊松过程,率分别为λ_1和λ_2,那么两个过程的到达组合形成一个率为λ_1 + λ_2的泊松过程。这对任意数量的过程都成立。
3. 它们以与其率成比例的概率赢得比赛。 在具有率λ_1的泊松过程的下一次到达与具有率λ_2的泊松过程的下一次到达之间的比赛中,第一个过程赢得比赛的概率是λ_1 / (λ_1 + λ_2)。

现在我们知道了泊松过程的基础知识,还需要一个小知识:

• 均匀-指数桥。 如果u是0到1之间均匀分布的随机变量,那么-ln(u)/λ就是具有率λ的指数分布。

考虑到均匀-指数桥,我们可以开始看到当算法给每一行分配一个键并按该键对总体进行排序时在做什么。它在所有总体行中进行一场比赛!在这场比赛中,每一行到达终点的时间是一个指数分布的随机变量,其率对应于该行的权重w_i。前n个到达构成样本。

为了证明这场比赛确实进行了我们想要的采样,我们将证明它等同于一系列单行抽取,每次抽取对于抽取时剩余的总体都是公平的。设总体的总权重为W,考虑一个任意权重为w_i的行i。算法将为其分配一个具有率w_i的指数分布随机变量,这对应于具有相同率的泊松过程的下一次到达。

现在考虑除i以外的所有行。它们也对应于泊松过程,其率等于它们的权重。我们可以将它们合并成一个率为W - w_i的组合过程。

现在,使用关于泊松比赛的规则,我们知道行i(由一个率为w_i的过程表示)将以概率

w_i / (w_i + (W - w_i)) = w_i / W

赢得与其他行(由一个率为W - w_i的组合过程表示)的比赛。

由于我们任意选择了行i,同样的论证适用于所有行。因此,每一行被抽中的概率等于其权重与总体总权重的比例。这证明了运行"指数比赛"让我们可以从总体中进行一次公平抽取。

但是,在我们抽取了一行之后,剩下的不就是一个稍小的新总体吗?我们不能在这个稍小的新总体上运行一场新的比赛来正确地抽取另一行吗?

我们可以。而且,由于泊松过程是无记忆的,我们不需要生成新的到达时间来运行这场新比赛。我们可以重用现有的到达时间,因为已经发生的到达对后续到达没有影响。因此,使用剩余到达时间抽取的下一行将是另一次公平抽取。

我们可以重复这个论证,以证明连续的行是相对于每次抽取时剩余的总体公平选择的。因此,A-ES算法通过连续抽取选择大小为n的样本,每次抽取相对于其剩余总体都是公平的。这就是证明。

更快采样的技巧

大多数大型分析数据集都会存储在面向列的存储格式中,如Parquet。从这种数据集读取时,你通常只需为你读取的列付费。(所谓"付费",我指的是等待查询引擎完成工作,但如果你在某些科技公司的云上运行查询,你可能真的要付费。)

例如,如果你的数据集包含一个有100列的表,但你只需要其中的4列,查询引擎通常只会读取那4列。相比之下,在面向行的数据存储中,即使你只想要每行中100个值中的4个,你通常也必须解码整行。此外,大多数面向列的存储都支持某种形式的过滤下推,允许存储引擎在过滤表达式评估为false时跳过行。这两个属性 - 按读取付费和过滤下推 - 是我们在采样时可以利用的。

假设我们有一个包含数十亿行和大约100列的Population表。我们如何高效地从中抽取1000行的加权样本?

我们可以使用前面讨论过的基本采样公式:

SELECT *
FROM Population
ORDER BY -LN(RANDOM()) / weight
LIMIT 1000

但是想想查询引擎必须做什么来执行这个查询。它必须读取和解码所有100列的数十亿行,以便将这些行传入排序/限制逻辑(通常实现为TOP_N操作)以确定要保留哪些行作为样本。即使样本只会保留这些行的0.00001%,你也必须付费读取整个Population表!

一个更快的方法是只读取我们需要的列来确定哪些行在样本中。假设我们的表有一个唯一标识每行的主键pk。以下是我们采样公式的一个变体,它只返回识别样本中行所需的主键:

SELECT pk
FROM Population
ORDER BY -LN(RANDOM()) / weight
LIMIT 1000

这个变体只强制查询引擎读取两列:pk和weight:是的,它仍然必须读取表中数十亿行的这两列,但这些列包含小值,可以快速扫描。毕竟,这正是面向列的存储设计擅长的。关键是我们不用付费读取大约100个额外的列,这些列的值我们99.99999%的时间都要丢弃。

一旦我们确定了样本中的行,我们就可以运行第二个查询来获取这些行的完整所需列集。

添加确定性

我们的采样算法依赖于随机化。如果我们用相同的输入运行算法两次,每次都会得到不同的结果。通常,这种非确定性正是我们想要的。

但有时不是。有时,能够控制算法依赖的骰子投掷是有用的。例如,有时能够重复一个样本是有用的。或者几乎重复一个样本。

为了让我们能够控制采样时使用的随机化性质,我们必须用确定性的伪随机函数替换对RANDOM的调用。一种常见的方法是对主键进行哈希,然后将哈希值映射到[0,1)范围内的数字。以下DuckDB宏pseudorandom_uniform就可以做到这一点:

CREATE MACRO pseudorandom_uniform(key, seed, index) AS (
  (hash(concat(key, seed, index)) & 0x7FFFFFFF) / cast(0x7FFFFFFF as DOUBLE)
);

我们可以改变seed和index参数,为相同的key生成独立的随机值。例如,如果我将seed固定为"demo-seed-20240601",并为key"key123"在index值1到10上生成随机数,我得到10个新的随机数:

SELECT i, pseudorandom_uniform('key123', 'demo-seed-20240601', i) AS u_key123
FROM generate_series(1, 10) AS t(i)

"key123"键和"demo-seed-20240601"种子的十个随机数。

要进行确定性采样,我们只需用对我们的函数pseudorandom_uniform()的调用替换对RANDOM()的调用。

现在我们可以进行确定性采样了,我们可以做更多有用的事情!例如,我们可以进行有放回采样。

有放回采样

早些时候,我们证明了A-ES算法允许我们进行无放回采样,作为一系列连续抽取,每次抽取从总体中移除一个项目,每次抽取相对于抽取时剩余的总体都是公平的。但如果我们想进行有放回采样呢?有放回采样要求我们在选择每个项目后将其返回总体,这样每次选择都相对于原始总体是公平的,并且个别项目可能被多次选择。

我们能在SQL中高效地实现有放回采样吗?是的!但这有点棘手。(我没有在任何地方找到这个算法的发表版本;如果你知道,请告诉我。我花了一些努力创建它,但我不会感到惊讶如果它已经为人所知。)

回想一下我们对A-ES算法的正确性证明。对于每一行具有权重w的行,该算法想象了一个相应的泊松过程,其速率为w,并用该过程的下一次到达来表示该行。该到达将发生在时间t = -ln(u)/w,其中u是范围(0,1)内均匀分布的随机数。然后算法按t值对所有行进行排序,并将前k次到达作为样本。

通过对这个算法进行一个小调整,我们可以进行有放回抽样。这个调整就是不仅考虑每行泊松过程的下一次到达,而是考虑所有到达。让t[i,j]表示第i行过程的第j次到达。由于我们知道在泊松过程中,连续到达之间的时间是指数分布的随机变量,我们可以对到达间隔时间进行累加,得到所需的到达时间。也就是说,t[i,j] = Σ[k=1 to j] -ln(u[i,k])/w[i],其中u[i,k]表示第i行的第k个均匀分布随机变量。

这个调整后的算法有一个小问题,那就是泊松过程理论上会产生无限多次到达。为每一行创建无限序列然后对所有这些序列的到达进行排序是不可行的。

幸运的是,我们可以避免这个问题!想想无放回抽样的A-ES算法与我们提出的有放回抽样的不可行算法之间的关系。我们可以用有放回算法来描述无放回算法:准备进行有放回抽样,但忽略j > 1的所有到达;剩下的到达必须是t[i,1]的形式,其中i表示相应的行。然后像以前一样,取这些剩余到达中的前k个作为样本。

现在考虑反过来,从无放回样本出发构造相应的有放回样本。设S是无放回抽样的行集。我们知道这些行由一组相应的到达时间t[i,1] (i ∈ S)表示。我们还知道,如果是有放回抽样,比赛会包括一些额外的到达时间t[i,j] (j > 1),这些时间可能会取代S中的一些获胜行。但关键是,我们还知道,如果R表示相应有放回样本中的行集,那么R必须包含在S中。这一说法源于这样一个事实:如果t[i,j] (j > 1)确实取代了前k次到达中的某一行,那么j > 1的要求意味着取代的行是某个在t[i,1]时更早到达样本的行的重复;因此,取代不会引入S以外的新行。

因此,如果我们有一个无放回样本,我们可以从它的行构造一个有放回样本。我们可以忽略总体中的所有其他行。这使问题变得更容易处理。

所以现在我们可以看到一个算法的雏形,用于进行大小为k的有放回抽样:

1. 首先,使用高效的A-ES算法进行k大小的无放回抽样。
2. 对于S中的每个抽样行,使用生成t[i,1]时相同的伪随机宇宙生成t[i,j] (j > 1)。
3. 取前k次到达作为样本。

你可能注意到,这个算法的第2步要求我们为S中的每一行创建无限多次到达。因此这一步需要O(k∞)的时间。当k很大时,这个时间可能会很长。

幸运的是,我们可以利用概率论将这个成本降低到O(k log k)。其思想是,如果行i预期在样本中出现μ次,那么当k很大时,它出现超过μ + c√μ次的可能性很小。所以我们不需要为S中的每一行生成完整的无限次到达;我们可以只生成μ + c√μ次到达,其中c足够大以满足我们个人的偏执程度。

这里是一个DuckDB宏的示例实现:

CREATE OR REPLACE MACRO sample_with_replacement(tbl, weight_col, size) AS (
    WITH RECURSIVE
    sample_without AS (
        SELECT *,
               -ln(random())/CAST("$weight_col" AS DOUBLE) AS t
        FROM "$tbl"
        ORDER BY t
        LIMIT "$size"
    ),
    arrivals(row_num, t) AS (
        SELECT row_number() OVER () AS row_num,
               t
        FROM sample_without
        UNION ALL
        SELECT row_num,
               t - ln(random())/CAST("$weight_col" AS DOUBLE)
        FROM arrivals
        INNER JOIN sample_without ON arrivals.row_num = sample_without.row_number
        WHERE row_num <= "$size" * 2 + 3 * sqrt("$size" * 2)
    )
    SELECT s.*
    FROM arrivals a
    INNER JOIN sample_without s ON a.row_num = s.row_number
    ORDER BY a.t
    LIMIT "$size"
);

有放回抽样的例子

作为一个我们可能想要进行有放回抽样而不是无放回抽样的例子,考虑以下代表抛一个有偏硬币可能结果的总体表ThreeToOne:

ThreeToOne总体表。

对于这个有偏硬币,"正面"出现的可能性是"反面"的3倍。我们可以通过从ThreeToOne总体表中有放回抽取10行来模拟抛这个硬币10次:

SELECT sample_with_replacement('ThreeToOne', 'weight', 10).*
FROM (VALUES (1)) _;

从ThreeToOne中有放回抽取10行样本的结果。

在这个样本中,我们得到了7个正面和3个反面。平均而言,我们预期在每个大小为10的样本中约有7.5个"正面",所以我们观察到的样本接近我们的预期。

但也许我们只是运气好。作为对我们SQL抽样逻辑的更强有力的测试,让我们取10,000个大小为40的样本,并查看所有样本中"正面"的计数。我们预期这个计数应该有一个二项分布(Binomial(size = 40, p = 3/4))。为了比较我们观察到的结果和预期分布,我将计算结果的经验分布并将其绘制在预期分布上。如你所见,观察到的分布与预期分布非常接近:

结论

抽样是一个强大的工具。通过我们刚刚讨论的SQL逻辑,你可以从几乎任何数据集中快速、轻松地抽取样本,无论数据集有多大。而且你可以进行有放回或无放回抽样。

使这一切成为可能的是与泊松过程理论的巧妙联系。这些过程是无记忆的和可合并的,它们的到达按照各自的速率赢得比赛。这些属性正是我们进行抽样比赛所需要的。

除了我们在本文中讨论的内容,我们还可以进一步利用这些属性。例如,作为性能优化,我们可以预测样本中最后一行的到达时间。然后我们可以用一个下推过滤器来增强我们的SQL抽样逻辑,该过滤器会消除到达时间大于t[k]的总体行,其中t[k]是某个常数。这种过滤发生在ORDER/LIMIT处理之前,可以大大加快查询速度,因为它在系统支持"延迟物化"的情况下,在完全读取行之前就消除了超过99.99%的行。

但这篇文章已经太长了,所以我现在就到此为止。

参考文献

Pavlos S. Efraimidis and Paul G. Spirakis. Weighted random sampling with a reservoir. Information Processing Letters, 97(5):181–185, 2006.