如何在Polars DataFrame间执行基于不等式条件的连接？

需要在两个Polars DataFrame之间执行基于m.date >= o.date的左连接，得到与指定SQL语句完全等效的结果。

给定的DataFrame #

from datetime import date
import polars as pl

stock_market_value = pl.DataFrame(
    {
        "date": [date(2022, 1, 1), date(2022, 2, 1), date(2022, 3, 1)],
        "price": [10.00, 12.00, 14.00]
    }
)

my_stock_orders = pl.DataFrame(
    {
        "date": [date(2022, 1, 15), date(2022, 2, 15)],
        "quantity": [2, 5] 
    }
)

需求对应的SQL语句 #

SELECT m.date, m.price * o.quantity AS portfolio_value
FROM stock_market_value m LEFT JOIN my_stock_orders o
  ON m.date >= o.date

期望输出 #

通过DuckDB执行等效查询得到的结果如下：

duckdb.sql("""
SELECT 
   m.date market_date, 
   o.date order_date, 
   price,
   quantity, 
   price * quantity AS portfolio_value
FROM stock_market_value m LEFT JOIN my_stock_orders o
  ON m.date >= o.date
""").pl()

输出结果：

shape: (4, 5)
┌─────────────┬────────────┬───────┬──────────┬─────────────────┐
│ market_date | order_date | price | quantity | portfolio_value │
│ ---         | ---        | ---   | ---      | ---             │
│ date        | date       | f64   | i64      | f64             │
╞═════════════╪════════════╪═══════╪══════════╪═════════════════╡
│ 2022-01-01  | null       | 10.0  | null     | null            │
│ 2022-02-01  | 2022-01-15 | 12.0  | 2        | 24.0            │
│ 2022-03-01  | 2022-01-15 | 14.0  | 2        | 28.0            │
│ 2022-03-01  | 2022-02-15 | 14.0  | 5        | 70.0            │
└─────────────┴────────────┴───────┴──────────┴─────────────────┘

为什么asof连接不适用 #

尝试过使用Polars的join_asof方法，但结果不符合预期：

Forward策略的asof连接 #

result_fwd = stock_market_value.join_asof(
    my_stock_orders, left_on="date", right_on="date", strategy="forward"
)

print(result_fwd)

输出：

shape: (3, 3)
┌────────────┬───────┬──────────┐
│ date       ┆ price ┆ quantity │
│ ---        ┆ ---   ┆ ---      │
│ date       ┆ f64   ┆ i64      │
╞════════════╪═══════╪══════════╡
│ 2022-01-01 ┆ 10.0  ┆ 2        │
│ 2022-02-01 ┆ 12.0  ┆ 5        │
│ 2022-03-01 ┆ 14.0  ┆ null     │
└────────────┴───────┴──────────┘

Backward策略的asof连接 #

result_bwd = stock_market_value.join_asof(
    my_stock_orders, left_on="date", right_on="date", strategy="backward"
)

print(result_bwd)

输出：

shape: (3, 3)
┌────────────┬───────┬──────────┐
│ date       ┆ price ┆ quantity │
│ ---        ┆ ---   ┆ ---      │
│ date       ┆ f64   ┆ i64      │
╞════════════╪═══════╪══════════╡
│ 2022-01-01 ┆ 10.0  ┆ null     │
│ 2022-02-01 ┆ 12.0  ┆ 2        │
│ 2022-03-01 ┆ 14.0  ┆ 5        │
└────────────┴───────┴──────────┘

可行解决方案 #

使用Polars的条件左连接，通过join方法的condition参数指定不等式匹配规则，即可实现与目标SQL完全一致的结果：

# 执行非等值左连接
result = stock_market_value.join(
    my_stock_orders,
    how="left",
    condition=pl.col("date", side="left") >= pl.col("date", side="right")
).rename(
    {"date": "market_date", "date_right": "order_date"}
).with_columns(
    (pl.col("price") * pl.col("quantity")).alias("portfolio_value")
)

print(result)

执行后输出：

shape: (4, 5)
┌─────────────┬────────────┬───────┬──────────┬─────────────────┐
│ market_date │ order_date │ price │ quantity │ portfolio_value │
│ ---         │ ---        │ ---   │ ---      │ ---             │
│ date        │ date       │ f64   │ i64      │ f64             │
╞═════════════╪════════════╪═══════╪══════════╪═════════════════╡
│ 2022-01-01  │ null       │ 10.0  │ null     │ null            │
│ 2022-02-01  │ 2022-01-15 │ 12.0  │ 2        │ 24.0            │
│ 2022-03-01  │ 2022-01-15 │ 14.0  │ 2        │ 28.0            │
│ 2022-03-01  │ 2022-02-15 │ 14.0  │ 5        │ 70.0            │
└─────────────┴────────────┴───────┴──────────┴─────────────────┘

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