如何在Python Polars中实现Pandas SQL式的多条件自连接？

处理大型DataFrame时遇到内存错误，计划用Python Polars替代Pandas SQL，但无法编写包含两个条件的自连接语法。

数据结构示例：

已实现的Pandas SQL版本自连接代码：

import pandas as pd
import datetime as dt
import pandasql as ps

d = {'Key':
          [1234,    1234,   1234,   1234,   1234,   1234,   1234,   
1234,   1234,   1234,   1234,   2456,   2456,   2456,   2456,   2456,   2456,   2456,   2456,   2456,   2456,   2456,   3754,   3754,   3754,   3754,   3754,   3754,   3754,   3754,   3754,   3754,   3754],
         'Field':[
              "Plumb",  "Plumb",    "Pear", "Plumb",    "Orange",   "Pear", "Plumb",    "Plumb",    "Pear", "Apple",    "Plumb",    "Orange",   "Orange",   "Apple",    "Apple",    "Pear", "Apple",    "Plumb",    "Plumb",    "Orange",   "Orange",   "Pear", "Plumb",    "Pear", "Plumb",    "Pear", "Apple",    "Plumb",    "Orange",   "Pear", "Apple",    "Pear", "Apple"],
         'DateColumn':[
              '2020-02-01', '2020-03-01',   '2020-04-01',   '2020-05-01',   '2020-06-01',   '2020-07-01',   '2020-08-01',   '2020-09-01',   '2020-10-01',   '2020-11-01',   '2020-12-01',   '2020-02-01',   '2020-03-01',   '2020-04-01',   '2020-05-01',   '2020-06-01',   '2020-07-01',   '2020-08-01',   '2020-09-01',   '2020-10-01',   '2020-11-01',   '2020-12-01',   '2020-02-01',   '2020-03-01',   '2020-04-01',   '2020-05-01',   '2020-06-01',   '2020-07-01',   '2020-08-01',   '2020-09-01',   '2020-10-01',   '2020-11-01',   '2020-12-01'

     ]}

df = pd.DataFrame(data=d)

df['DateColumn'] = pd.to_datetime(df['DateColumn'])
df['PreviousMonth'] = df['DateColumn'] - pd.DateOffset(months=1)


df_output = ps.sqldf("""
select
a.Key
,a.Field
,b.Field as PreviousField
,a.DateColumn
,b.DateColumn as PreviousDate

from df as a, df as b
where a.Key = b.Key
and b.DateColumn = a.PreviousMonth

""")

print(df_output.head())

该代码输出：

Key   Field                  DateColumn                PreviousDate
0  1234   Plumb  2020-03-01 00:00:00.000000  2020-02-01 00:00:00.000000
1  1234    Pear  2020-04-01 00:00:00.000000  2020-03-01 00:00:00.000000
2  1234   Plumb  2020-05-01 00:00:00.000000  2020-04-01 00:00:00.000000
3  1234  Orange  2020-06-01 00:00:00.000000  2020-05-01 00:00:00.000000
4  1234    Pear  2020-07-01 00:00:00.000000  2020-06-01 00:00:00.000000

尝试了Polars的join语句但不知道添加两个连接条件：

data_output = df.join(df, left_on='Key', right_on='Key')

在Polars中，实现多条件自连接可以通过以下两种方式：

方式1：使用join方法的多列连接条件 #

直接将多个连接条件以列表形式传入left_on和right_on参数，同时需要给右侧DataFrame重命名列避免冲突：

import polars as pl

# 转换为Polars DataFrame并处理日期列
pl_df = pl.DataFrame(d)
pl_df = pl_df.with_columns(
    pl.col("DateColumn").str.to_date(),
    PreviousMonth=pl.col("DateColumn").str.to_date().dt.offset_by("-1mo")
)

# 自连接：匹配Key和日期条件
output = pl_df.join(
    pl_df.rename({"Field": "PreviousField", "DateColumn": "PreviousDate"}),
    left_on=["Key", "PreviousMonth"],
    right_on=["Key", "PreviousDate"],
    how="inner"
).select(
    "Key",
    "Field",
    "PreviousField",
    "DateColumn",
    "PreviousDate"
)

print(output.head())

方式2：使用join配合condition参数（Polars >= 0.19.0） #

如果需要更灵活的连接条件，可以使用condition参数指定布尔表达式：

import polars as pl

# 转换为Polars DataFrame并处理日期列
pl_df = pl.DataFrame(d)
pl_df = pl_df.with_columns(
    pl.col("DateColumn").str.to_date(),
    PreviousMonth=pl.col("DateColumn").str.to_date().dt.offset_by("-1mo")
)

output = pl_df.join(
    pl_df.rename({"Field": "PreviousField", "DateColumn": "PreviousDate"}),
    on="Key",
    condition=pl.col("PreviousMonth") == pl.col("PreviousDate"),
    how="inner"
).select(
    "Key",
    "Field",
    "PreviousField",
    "DateColumn",
    "PreviousDate"
)

print(output.head())

两种方式都会得到和Pandas SQL版本一致的输出，且Polars的内存效率更高，适合处理大型数据集。

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