DeepSeek总结的DuckDB使用 WITH RECURSIVE 和 USING KEY 进行聚合的特性

原文地址：https://github.com/duckdb/duckdb/commit/6b441dd39844b1b303c403623452e2810ef0e195

特性：在 USING KEY 中使用聚合 (#19481)

大家好，

TLDR（摘要）：

此 PR 为 USING KEY 特性引入了新功能。简而言之，我们泛化了 USING KEY 的原则，通过让用户更精确地控制非键列的处理方式。他们可以选择性地为有效载荷（即非键）列定义自己的聚合函数。这使得用户能够保留在迭代中产生的、并非最后一个值的其他值。例如，他们可以保存列中的最小值（min），或者依赖于该行中其他列的值（arg_min）。

得益于这种新的泛化，USING KEY 可以用于进一步简化算法，例如最短路径算法，这已被他人提及。

# 实现方式和语法是如何工作的？

USING KEY 的实现使用了一个聚合哈希表。该表包含键列和有效载荷列，哈希表将在有效载荷列上执行聚合函数。

在此之前，我们使用 last 聚合函数来实现 USING KEY 特有的 upsert 行为：如果某次迭代产生的键已存在于字典中，则最近一次 迭代会替换该键对应的有效载荷。通过此 PR，用户可以使用替代的聚合函数（例如 min、max、avg 等），从而覆盖默认的 last。

以下查询仅为解释 USING KEY 执行行为而进行了简化。

现在，我们可以在 USING KEY 子句中为特定的有效载荷列指定额外的聚合函数。我们可以通过列引用（例如 max(<column_ref>)）在聚合中引用输出列。

sql 复制代码

 WITH RECURSIVE tbl(a, b) USING KEY (a, max(b)) AS (
      SELECT 1, 5
          UNION
      SELECT a, b - 1
      FROM tbl
      WHERE b > 0
) TABLE tbl;
┌───────┬───────┐
│   a   │   b   │
│ int32 │ int32 │
├───────┼───────┤
│   1   │   5   │
└───────┴───────┘

如查询结果所示，此查询实际上会终止，而不是在每次迭代中产生行 (1, 5)：列 b 中的 max 聚合值被维护（并且在查询使用 recurring.tbl 时可被获取），而工作表（将包含行 (1, 5)，然后是 (1, 4)，最后是 (1, 0)）的行为继续遵循常规递归 CTE 的语义。

值得注意的是，在循环表中，我们可以访问截至目前计算出的聚合值：我们可以检查到递归中此点为止的聚合结果。（这实际上允许我们将循环表用作滑动窗口，来计算累加和等。）

sql 复制代码

 WITH RECURSIVE tbl(a, b, c) USING KEY (a, avg(b), list(c)) AS (
      SELECT 1, 5, NULL :: DOUBLE
          UNION
      SELECT a.a, a.b - 1, b.b
      FROM tbl AS a, recurring.tbl AS b
      WHERE a.b > 0
) TABLE tbl;
┌───────┬────────┬─────────────────────────────────┐
│   a   │   b    │               c                 │
│ int32 │ double │            double[]             │
├───────┼────────┼─────────────────────────────────┤
│   1   │  2.5   │ [NULL, 5.0, 4.5, 4.0, 3.5, 3.0] │
└───────┴────────┴─────────────────────────────────┘

然而，也可以在聚合中使用任意表达式，而不仅仅是纯列引用。唯一的限制是，任何包含的列引用必须是递归 CTE 模式的一部分。

在这种情况下，我们还需要指定将保存聚合结果的输出列的名称。我们可以通过使用命名参数语法 b := max(a+b) 或 max(a+b) AS b 来实现。这两种变体都可以。

sql 复制代码

 WITH RECURSIVE tbl(a, b) USING KEY (a, max(b*3) AS b) AS (
      SELECT 1, 5
          UNION
      SELECT a, b - 1
      FROM tbl
      WHERE b > 0
) TABLE tbl;
┌───────┬───────┐
│   a   │   b   │
│ int32 │ int32 │
├───────┼───────┤
│   1   │   15  │
└───────┴───────┘

# 新的方法

另一个积极的副作用是，当使用非增量聚合函数（如 median）时，我们可以节省内存。对于这些聚合函数，我们需要来自所有迭代的所有计算值。对于普通的递归 CTE 和当前 USING KEY 的实现，有两种方式可以使用非增量聚合函数。

旧方式 ：像以前使用普通递归 CTE 一样，我们在联合表中收集所有必要的行，然后在 CTE 计算完成后使用带有 GROUP BY 的聚合函数。

sql 复制代码

WITH RECURSIVE tbl(a, b, payload) AS (
	SELECT 1, 1, a
	FROM (SELECT list(a) FROM generate_series(1,5000) AS _(a))_(a)
		UNION
	SELECT a, b+1, payload
	FROM tbl WHERE b < 200000)
SELECT median(b)
FROM tbl
GROUP BY a;

这种方法的缺点是它复制了整个有效载荷。如果有效载荷的某些部分与聚合函数完全无关，我们就在不必要地增加内存占用。

USING KEY 方式 ：旧的 USING KEY 语义只记住最近产生的行，因此我们必须添加一个新列来存储聚合所需的所有值。

sql 复制代码

WITH RECURSIVE tbl(a, b, acc, payload) USING KEY (a) AS (
	SELECT 1, 1, [1],  a
	FROM (SELECT list(a) FROM generate_series(1,5000) AS _(a))_(a)
		UNION
	SELECT a, b+1, acc || [b+1], payload
	FROM tbl WHERE b < 200000)
SELECT a, list_median(b) FROM tbl;

我们没有重复的有效载荷，但仍然需要添加一个额外的列来计算非增量聚合，以及一个额外的数组和对其的操作来模拟"记忆"部分。

使用 USING KEY 进行聚合的方式：通过我在此处实现的新特性，我们既可以获得去重有效载荷的优势，也无需添加和管理新列。

sql 复制代码

WITH RECURSIVE tbl(a, b, payload) USING KEY (a, median(b)) AS (
	SELECT 1, 1, a
	FROM (SELECT list(a) FROM generate_series(1,5000) AS _(a))_(a)
		UNION
	SELECT a, b+1, payload
	FROM tbl WHERE b < 200000)
SELECT a, b FROM tbl

在这里，我们将聚合函数添加到 USING KEY 子句中。聚合函数本身会处理值的存储方式，确保有效载荷不被重复。