DuckDB中秩相关性计算性能远低于常规相关性的原因排查

两段仅差Rank计算的DuckDB代码，第二段因先计算rank，性能比第一段慢约5倍。仅多了Rank相关计算却有如此大的性能差距，怀疑是查询实现存在低效问题，同时对DuckDB的查询优化能力存疑。注：用Polars执行类似计算不会出现如此显著的性能下降。

测试代码 #

Pearson相关计算（耗时~19.46秒） #

import time

import duckdb
import numpy as np
import polars as pl

# 生成示例数据
rng = np.random.default_rng(1)
nrows = 5_000_000
df = pl.DataFrame(
    dict(
        id=rng.integers(1, 1_000, nrows),
        id2=rng.integers(1, 10, nrows),
        id3=rng.integers(1, 500, nrows),
        v1=rng.normal(0, 1, nrows),
        v2=rng.normal(0, 1, nrows),
    )
)

start = time.perf_counter()
res = duckdb.sql(
    """
    WITH cte AS (
        SELECT df.id2, df.id3, df2.id3 AS id3_right,
            df.v1, df2.v1 AS v1_right,
            df.v2, df2.v2 AS v2_right
        FROM df
        JOIN df AS df2
        ON (
            df.id = df2.id
            AND df.id2 = df2.id2
            AND df.id3 > df2.id3
            AND df.id3 < df2.id3 + 30
        )
    )
    SELECT id2, id3, id3_right, corr(v1, v1_right) AS v1, corr(v2, v2_right) AS v2
    FROM cte
    GROUP BY id2, id3, id3_right
    """
).pl()
print(time.perf_counter() - start)
# 输出：19.462523670867085

Rank相关计算（耗时~104.54秒） #

start = time.perf_counter()
res2 = duckdb.sql(
    """
    WITH cte AS (
        SELECT df.id2, df.id3, df2.id3 AS id3_right,
            RANK() OVER (g ORDER BY df.v1) AS v1,
            RANK() OVER (g ORDER BY df2.v1) AS v1_right,
            RANK() OVER (g ORDER BY df.v2) AS v2,
            RANK() OVER (g ORDER BY df2.v2) AS v2_right
        FROM df
        JOIN df AS df2
        ON (
            df.id = df2.id
            AND df.id2 = df2.id2
            AND df.id3 > df2.id3
            AND df.id3 < df2.id3 + 30
        )
        WINDOW g AS (PARTITION BY df.id2, df.id3, df2.id3)
    )
    SELECT id2, id3, id3_right, corr(v1, v1_right) AS v1, corr(v2, v2_right) AS v2
    FROM cte
    GROUP BY id2, id3, id3_right
    """
).pl()
print(time.perf_counter() - start)
# 输出：104.54312287131324

核心原因分析 #

• 窗口分区与聚合分组重复：窗口g的分区键和后续GROUP BY的分组键完全一致，导致DuckDB先对每个分区做排序计算Rank，后续GROUP BY又重新聚合，产生重复计算开销。
• Join后数据量膨胀：自Join加上id3的范围条件会生成远超原始数据量的中间数据集，在这个大数据集上执行4次独立的窗口Rank计算，每个都要做分区排序，性能开销被大幅放大。
• 引擎优化差异：Polars的执行引擎对“窗口计算+聚合”的组合逻辑做了更智能的融合优化，避免了重复的分区和排序操作，因此性能下降不明显。

优化方案 #

1. 提前计算Rank，再做Join：先在原始数据上按id2, id3分区计算Rank，再执行Join操作，大幅减少窗口计算的数据量。调整后的查询示例：

   WITH ranked_df AS (
       SELECT 
           id, id2, id3,
           RANK() OVER (PARTITION BY id2, id3 ORDER BY v1) AS v1_rank,
           RANK() OVER (PARTITION BY id2, id3 ORDER BY v2) AS v2_rank
       FROM df
   )
   SELECT 
       rd.id2, rd.id3, rd2.id3 AS id3_right,
       corr(rd.v1_rank, rd2.v1_rank) AS v1,
       corr(rd.v2_rank, rd2.v2_rank) AS v2
   FROM ranked_df rd
   JOIN ranked_df rd2
   ON (
       rd.id = rd2.id
       AND rd.id2 = rd2.id2
       AND rd.id3 > rd2.id3
       AND rd.id3 < rd2.id3 + 30
   )
   GROUP BY rd.id2, rd.id3, rd2.id3
   

2. 升级DuckDB版本：新版本的DuckDB通常会优化窗口函数与聚合的组合逻辑，减少不必要的计算开销。
3. 启用窗口优化：如果使用的DuckDB版本支持，可执行SET optimizer_enable_window_optimizations = true;，让优化器尝试合并窗口与聚合操作。

内容的提问来源于stack exchange，提问作者lebesgue