在Polars数据处理场景中,按前缀合并同类型列并保留非匹配列、对齐长度是常见的需求,比如处理带时间戳前缀的指标列、分组统计后的结果列等,都可以通过Polars的API组合实现。

环境准备
首先需要确保已经安装Polars库,如果未安装可以通过pip命令安装:
# 安装Polars库 pip install polars
导入需要的模块并创建测试数据集,方便后续演示操作逻辑:
import polars as pl
# 创建测试DataFrame,包含不同前缀的列和非匹配列
df = pl.DataFrame({
"user_id": [1, 2, 3],
"score_math": [90, 85, 92],
"score_english": [88, 91, 87],
"score_physics": [95, 89, 93],
"extra_info": ["A", "B", "C"]
})
print("原始数据:")
print(df)
核心实现思路
实现按前缀合并列并复制非匹配列对齐长度的核心步骤如下:
- 第一步:识别所有需要合并的带前缀的列,提取对应的前缀和列名
- 第二步:对每个前缀下的列进行合并操作,生成新的合并列
- 第三步:筛选出不带目标前缀的非匹配列,保留这些列的原始数据
- 第四步:将合并后的新列和非匹配列组合,确保所有列长度一致
完整实现代码
以下是完整的实现函数,支持指定前缀列表,自动合并对应列并保留非匹配列:
def merge_columns_by_prefix(df: pl.DataFrame, prefixes: list[str]) -> pl.DataFrame:
"""
按前缀合并列,并保留非匹配列对齐长度
:param df: 输入的Polars DataFrame
:param prefixes: 需要合并的列前缀列表
:return: 合并后的DataFrame
"""
merged_cols = []
used_cols = set()
# 处理每个前缀的合并逻辑
for prefix in prefixes:
# 筛选当前前缀对应的所有列
prefix_cols = [col for col in df.columns if col.startswith(prefix)]
if not prefix_cols:
continue
# 合并同前缀的列,这里使用水平拼接,自动对齐长度
merged_series = pl.concat([df[col] for col in prefix_cols], how="horizontal")
# 生成新的合并列名
new_col_name = f"{prefix}_merged"
merged_cols.append(merged_series.alias(new_col_name))
# 记录已经使用的列
used_cols.update(prefix_cols)
# 筛选非匹配列,即不在任何前缀列表中的列
non_matched_cols = [col for col in df.columns if col not in used_cols]
non_matched_series = [df[col] for col in non_matched_cols]
# 组合合并列和非匹配列
result_df = pl.DataFrame(merged_cols + non_matched_series)
return result_df
# 调用函数,指定需要合并的前缀为score_
result = merge_columns_by_prefix(df, ["score_"])
print("合并后的数据:")
print(result)
代码逻辑解析
上述函数的核心逻辑分为几个部分:
前缀列筛选
通过startswith方法筛选每个前缀对应的列,避免手动指定列名,提升函数的通用性。
列合并操作
使用pl.concat方法并设置how="horizontal",可以将同前缀的多个列水平拼接,Polars会自动处理长度对齐的问题,如果列长度不一致会按照最长列的长度复制短列的内容。
非匹配列保留
通过对比所有列名和使用过的前缀列名,得到非匹配列列表,直接将这些列的原始数据加入结果集,保证非匹配列不会被丢弃。
多场景示例
场景1:多个不同前缀合并
如果数据中有多个不同的前缀需要合并,可以直接在前缀列表中添加多个前缀:
# 创建包含多个前缀的测试数据
df2 = pl.DataFrame({
"id": [1, 2],
"temp_2023": [20, 22],
"temp_2024": [21, 23],
"hum_2023": [60, 62],
"hum_2024": [61, 63],
"location": ["北京", "上海"]
})
# 合并temp_和hum_前缀的列
result2 = merge_columns_by_prefix(df2, ["temp_", "hum_"])
print("多前缀合并结果:")
print(result2)
场景2:部分前缀无匹配列
如果指定的前缀在数据中不存在对应列,函数会自动跳过该前缀,不会报错:
# 指定一个不存在的前缀
result3 = merge_columns_by_prefix(df, ["score_", "grade_"])
print("包含不存在前缀的合并结果:")
print(result3)
注意事项
- 合并的同前缀列需要是相同数据类型,否则Polars会报错,提前做类型校验可以避免问题
- 如果非匹配列的长度和合并列长度不一致,Polars会自动按照最长列的长度复制非匹配列的内容,保证所有列长度一致
- 合并后的新列名可以自定义,上述示例中默认使用前缀加_merged的命名方式,可根据需求修改