为何Polars中Join前收集LazyFrame可解决索引不一致问题？

可复现示例 #

以下代码可复现问题：初始LazyFrame包含平面中点的两两距离，仅保留唯一距离对后，构建含唯一标签的labels LazyFrame。直接关联LazyFrame时，同索引间的距离（如A-A、B-B）结果不符合预期，提前收集为DataFrame后结果正确。

import polars as pl
import numpy as np

data = {
    "col_1": ["A", "A", "A", "B", "B", "B", "C", "C", "C"],
    "col_2": ["B", "C", "A", "A", "C", "B", "A", "B", "C"],
    "col_3": [1.0, 2.0, 0.0, 1.0, 1.5, 0.0, 2.0, 1.5, 0.0]
}

# 创建LazyFrame
pairwise_distances = pl.LazyFrame(data)

# 拼接并排序两个元素，返回带||的字符串
def concat_and_sort(a, b):
    return "||".join(sorted([a, b]))

# 添加配对列并保留唯一配对
pairwise_distances = pairwise_distances.with_columns(
    pl.struct(["col_1", "col_2"]).map_elements(
        lambda x: concat_and_sort(x["col_1"], x["col_2"]),
        return_dtype=pl.String
    ).alias("pairs")
).unique(subset="pairs").select(pl.col("col_1"), pl.col("col_2"), pl.col("col_3"))

# 拼接col_1和col_2，生成带索引的唯一标签表
labels = (
    pl.concat([
        pairwise_distances.select(pl.col("col_1").alias("label")),
        pairwise_distances.select(pl.col("col_2").alias("label"))
    ])
    .unique(keep="first")
    .with_row_count(name="index")
)

# 收集为DataFrame（用于对比）
labels_df = labels.collect()
pairwise_distances_df = pairwise_distances.collect()

# 直接用LazyFrame关联（结果错误）
data_joined_direct = pairwise_distances.join(
    labels,
    left_on="col_1",
    right_on="label",
    how="left"
).rename({"index": "index_1"})

data_joined_direct = data_joined_direct.join(
    labels,
    left_on="col_2",
    right_on="label",
    how="left"
).rename({"index": "index_2"})

result_direct_df = data_joined_direct.select(["index_1", "index_2", "col_3"]).collect()

# 收集后用DataFrame关联（结果正确）
data_joined_collected = pairwise_distances_df.join(
    labels_df,
    left_on="col_1",
    right_on="label",
).rename({"index": "index_1"})

data_joined_collected = data_joined_collected.join(
    labels_df,
    left_on="col_2",
    right_on="label",
).rename({"index": "index_2"})

result_collected = data_joined_collected.select(["index_1", "index_2", "col_3"])

问题描述 #

使用Polars处理数据集时，需从两列生成唯一标签，再通过Join操作获取标签对应的索引。但发现直接使用LazyFrame执行Join时索引结果错误，而将LazyFrame收集为DataFrame后再执行Join，索引则完全正确。

1. 创建labels LazyFrame #

import polars as pl

# 假设data是从文件读取的LazyFrame
data = pl.scan_csv(
    source=source_filepath,
    separator='\t',
    has_header=False,
)

# 拼接col_1和col_2，生成带索引的唯一标签表
labels = (
    pl.concat([
        data.select(pl.col("col_1").alias("label")),
        data.select(pl.col("col_2").alias("label"))
    ])
    .unique(keep="first")
    .with_row_count(name="label_index")
)

2. 未收集时执行Join（索引错误） #

# 关联获取index_1
data = data.join(
    labels,
    left_on="col_1",
    right_on="label",
).rename({"label_index": "index_1"})

# 关联获取index_2
data = data.join(
    labels,
    left_on="col_2",
    right_on="label",
).rename({"label_index": "index_2"})

result_df = data.select(["index_1", "index_2", "col_3"]).collect()

3. 收集后执行Join（索引正确） #

# 提前收集为DataFrame
labels_df = labels.collect()
data_df = data.collect()

# 关联获取index_1
data_df = data_df.join(
    labels_df,
    left_on="col_1",
    right_on="label",
    how="left"
).rename({"label_index": "index_1"})

# 关联获取index_2
data_df = data_df.join(
    labels_df,
    left_on="col_2",
    right_on="label",
    how="left"
).rename({"label_index": "index_2"})

result_df = data_df.select(["index_1", "index_2", "col_3"])

问题原因 #

Polars的LazyFrame是延迟计算机制，当你直接用labels这个LazyFrame多次执行Join时，Polars会每次重新执行labels的完整构建逻辑（包括concat、unique、with_row_count）。而unique(keep="first")的结果依赖于数据的输入顺序，每次重新执行时，LazyFrame的执行计划可能导致数据顺序变化，或者两次扫描原数据的顺序不一致，最终生成的label_index（行号）会出现差异，导致两次Join得到的索引不匹配。

而提前收集为DataFrame后，labels_df是固定在内存中的数据集，多次Join时使用的是同一个固定的索引映射，因此结果正确。

解决方法 #

无需提前收集LazyFrame即可保证结果正确的方案有三种：

方案1：缓存labels LazyFrame #

在构建labels后调用.cache()方法，让Polars只计算一次并缓存结果，后续Join都会复用缓存的版本：

labels = (
    pl.concat([
        data.select(pl.col("col_1").alias("label")),
        data.select(pl.col("col_2").alias("label"))
    ])
    .unique(keep="first")
    .with_row_count(name="label_index")
    .cache()  # 关键：缓存计算结果
)

方案2：用窗口函数直接生成索引 #

不需要单独构建labels表，直接在原数据中用dense_rank()窗口函数生成全局唯一索引，避免多次Join的问题：

# 生成全局唯一索引（默认从1开始）
data = data.with_columns(
    pl.col("col_1").rank(method="dense").alias("index_1"),
    pl.col("col_2").rank(method="dense").alias("index_2")
)

# 如需从0开始索引，执行减1操作
data = data.with_columns(
    (pl.col("index_1") - 1).alias("index_1"),
    (pl.col("index_2") - 1).alias("index_2")
)

result_df = data.select(["index_1", "index_2", "col_3"]).collect()

方案3：固定labels的数据源 #

将labels收集为DataFrame后再转回LazyFrame，本质是固定数据内容，避免重复执行构建逻辑：

labels = pl.LazyFrame(labels.collect())

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