Polars中多次连续表连接的更优（更快、内存更高效）实现方案问询

我最近在处理大型Polars数据集时，遇到了多次连续表连接速度拉胯的问题，于是自己摸索出了一个替代方案：先把需要合并的DataFrame做透视转换，只执行一次表连接，再结合pl.coalesce和条件逻辑来整理出目标数据。在数据量逐渐变大的场景下，这个方法确实比多次连续连接要快，但我还是想问问有没有更优（更快、内存效率更高）的实现方式？

实际业务场景中，我可能会针对更多特征、更多标的，以及规模更大的篮子数据重复这类操作。

测试数据准备 #

先给大家看看我用来测试的基础数据代码：

import polars as pl
from polars import col
from vega_datasets import data

df = pl.from_pandas(data.stocks())
list_symbols = df.select(col('symbol').unique()).to_series().to_list()
# 生成的标的列表：["IBM", "MSFT", "AAPL", "AMZN", "GOOG"]

basket = pl.DataFrame(
    {
        "MSFT": ["AMZN", "GOOG"],
        "AMZN": ["MSFT", "GOOG"],
        "GOOG": ["AAPL", "IBM"],
        "IBM": ["AMZN", "AAPL"],
        "AAPL": ["AMZN", "IBM"],
    }
).transpose(
    include_header=True,
    header_name="symbol",
    column_names=["symbol_1", "symbol_2"],
)

现有两种实现方案 #

方案1：多次连续左连接

这是我最开始的实现思路，通过多次左连接来关联对应标的的价格数据（实际场景中连接次数会远多于2次）：

(
    df
    .join(basket, on='symbol', how='left')
    .join(df.select('date','symbol', col('price').name.suffix('_1')), 
          left_on = ['date', 'symbol_1'],
          right_on = ['date', 'symbol'],
          how='left',
         )
    .join(df.select('date','symbol', col('price').name.suffix('_2')), 
          left_on = ['date', 'symbol_2'],
          right_on = ['date', 'symbol'],
          how='left',
         )
)

方案2：透视宽表+单次连接+coalesce条件提取

这是我自己想到的优化方案，先把价格数据透视成宽表，只做一次基于日期的连接，最后通过条件匹配提取对应标的的价格：

df_pivot = (
    # 日常工作中我主要用LazyFrame的透视实现，这里为了简化用了EagerFrame
    df.pivot(index='date', on='symbol', values='price')
)

(
    df.join(basket, on="symbol", how="left")
    # 核心：用一次连接替代多次连续连接
    .join(
        df_pivot.select('date', pl.exclude('date').name.suffix('_price_to_drop')),
        on="date",
        how="left",
    )
    .with_columns(
        *[
            pl.coalesce(
                pl.when(col(f'symbol_{i}')==symbol)
                .then(col(f'{symbol}_price_to_drop'))
                for symbol in list_symbols
            ).alias(f'price_{i}')
            for i in [1,2]
        ]
    )
    .select(pl.exclude("^.*to_drop$"))
)

备注：日常工作中我主要使用LazyFrame来处理数据，以上示例为了演示方便用了EagerFrame。

备注：内容来源于stack exchange，提问作者AD AD