如何在DuckDB的关系型API中使用CASE WHEN语句？

数据定义 #

data = {'id': [1, 1, 1, 2, 2, 2],
 'd': [1, 2, 3, 1, 2, 3],
 'sales': [1, 4, 2, 3, 1, 2]}

目标SQL查询 #

要转换的原始SQL查询如下（通过DuckDB执行）：

import duckdb
import polars as pl
df = pl.DataFrame(data)
duckdb.sql("""
select *, case when count(sales) over w then sum(sales) over w else null end as rolling_sales
from df
window w as (partition by id order by d rows between 1 preceding and current row)
""")

已完成的部分代码 #

目前已经实现了不带CASE WHEN逻辑的关系型API写法：

rel = duckdb.table("df")
rel.sum(
    "sales",
    projected_columns="*",
    window_spec="over (partition by id order by d rows between 1 preceding and current row) as rolling_sales",
)

解决方案 #

要在关系型API中加入CASE WHEN逻辑，直接用select方法结合表达式字符串即可，两种实现方式如下：

方式1：直接在select中写完整窗口表达式 #

import duckdb
import polars as pl

data = {'id': [1, 1, 1, 2, 2, 2],
 'd': [1, 2, 3, 1, 2, 3],
 'sales': [1, 4, 2, 3, 1, 2]}
df = pl.DataFrame(data)

rel = duckdb.table("df")
result = rel.select(
    "*",
    """case when count(sales) over (partition by id order by d rows between 1 preceding and current row) 
            then sum(sales) over (partition by id order by d rows between 1 preceding and current row) 
            else null end as rolling_sales"""
)

# 输出结果
print(result.show())

方式2：先定义窗口再引用（复用窗口逻辑） #

如果想避免重复写窗口定义，可以先用window方法声明窗口别名，再在select的CASE WHEN中引用：

rel = duckdb.table("df")
result = rel.window(
    w="partition by id order by d rows between 1 preceding and current row"
).select(
    "*",
    "case when count(sales) over w then sum(sales) over w else null end as rolling_sales"
)

这两种写法都能实现目标SQL的逻辑，同时保持关系型API的可读性。

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