Python+Polars：能否标记列作为匹配用的“主键”或“标签”列？

Polars本身没有内置的mark_primary这类方法，但可以通过以下两种便捷方式实现标记特定列为匹配用途列的需求：

方法一：借助DataFrame的自定义metadata+辅助函数 #

利用Polars DataFrame支持自定义metadata的特性，用辅助函数实现标记和获取操作：

import polars as pl

def mark_primaries(df: pl.DataFrame, cols: list[str]) -> pl.DataFrame:
    # 将指定列名存入DataFrame的metadata
    return df.with_metadata({"primaries": cols})

def get_primaries(df: pl.DataFrame) -> list[str]:
    # 从metadata中读取标记的列，默认返回空列表
    return df.metadata.get("primaries", [])

# 示例使用
# 构造初始数据
ids_a = pl.Series("id_a", [1, 2, 3])
ids_b = pl.Series("id_b", [10, 20, 30])
data_x = pl.Series("data_x", [100, 200, 300])
df = pl.DataFrame([ids_a, ids_b, data_x])

# 标记匹配用列
df = mark_primaries(df, ["id_a", "id_b"])

# 构造另一张表并标记
other_ids_a = pl.Series("id_a", [2, 3, 4])
other_ids_b = pl.Series("id_b", [20, 30, 40])
data_y = pl.Series("data_y", [2000, 3000, 4000])
other_df = pl.DataFrame([other_ids_a, other_ids_b, data_y])
other_df = mark_primaries(other_df, ["id_a", "id_b"])

# 获取交集作为join的on参数
join_on = set(get_primaries(df)).intersection(get_primaries(other_df))
joined_df = df.join(other_df, on=list(join_on))
print(joined_df)

方法二：自定义DataFrame子类 #

通过继承pl.DataFrame，添加专属方法，让调用方式更贴近你的伪代码：

import polars as pl

class PrimaryDataFrame(pl.DataFrame):
    def mark_primary(self, cols: list[str]) -> "PrimaryDataFrame":
        # 将标记列存入metadata
        return self.with_metadata({"primaries": cols})
    
    def primaries(self) -> list[str]:
        # 返回标记的匹配用列
        return self.metadata.get("primaries", [])

# 示例使用
ids_a = pl.Series("id_a", [1, 2, 3])
ids_b = pl.Series("id_b", [10, 20, 30])
data_x = pl.Series("data_x", [100, 200, 300])

df = PrimaryDataFrame([ids_a, ids_b, data_x])
df = df.mark_primary(["id_a", "id_b"])

other_ids_a = pl.Series("id_a", [2, 3, 4])
other_ids_b = pl.Series("id_b", [20, 30, 40])
data_y = pl.Series("data_y", [2000, 3000, 4000])
other_df = PrimaryDataFrame([other_ids_a, other_ids_b, data_y])
other_df = other_df.mark_primary(["id_a", "id_b"])

join_on = set(df.primaries()).intersection(other_df.primaries())
joined_df = df.join(other_df, on=list(join_on))
print(joined_df)

注意事项 #

• Polars的部分操作（如完全重构DataFrame的操作）可能会重置metadata，若需要在数据转换后仍保留标记，需重新执行标记操作或提前确认操作是否会保留metadata。
• 两种方案都支持后续在group_by时直接调用标记列：df.group_by(by=get_primaries(df))（方法一）或df.group_by(by=df.primaries())（方法二）。

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