在 Polars 中从字典合并 DataFrame 并保留来源名称

本教程详细介绍了如何在 polars 中高效地将存储在字典中的多个 dataframe 垂直合并为一个单一的 dataframe，同时为每个原始 dataframe 添加一个新列，记录其在字典中的键（即来源名称）。通过结合使用列表推导式、`with_columns` 和 `pl.lit`，以及 `pl.concat` 函数，我们能够优雅地解决在数据整合过程中保留数据来源信息的常见需求，从而实现类似 r tidyverse 中 `bind_rows(.id)` 的功能。

在数据分析和处理中，我们经常会遇到需要从多个来源（例如不同的 Excel 工作表、CSV 文件或数据库查询结果）导入数据，并将它们存储在一个字典中，其中每个键代表一个数据源的名称，对应的值则是一个 Polars DataFrame。随后的一个常见需求是将这些分散的 DataFrame 合并成一个统一的 DataFrame，同时希望在新合并的数据中能够追溯每一行数据来自哪个原始来源。

例如，当我们从一个多工作簿的 Excel 文件中读取数据时，polars.read_excel 通常会返回一个字典，其中键是工作表名称，值是对应的 DataFrame。此时，我们希望将这些 DataFrame 垂直堆叠起来，并添加一个新列来标识原始的工作表名称。这类似于 R Tidyverse 中 bind_rows(.id = "ID") 的便捷功能。

Polars 提供了 pl.concat 函数用于垂直或水平地连接 DataFrame。然而，直接对字典中的 DataFrame 列表使用 pl.concat 会导致丢失原始的字典键（即来源名称）。为了解决这个问题，我们需要在合并之前，为每个 DataFrame 添加一个包含其来源名称的新列。

合并 DataFrame 并保留来源名称

要实现这一目标，我们可以结合使用 Python 的列表推导式、Polars 的 with_columns 方法和 pl.lit 函数，最终通过 pl.concat 完成合并。

核心思路：

1. 遍历字典中的每个键值对（名称和 DataFrame）。
2. 对于每个 DataFrame，使用 with_columns 方法添加一个新的列。
3. 新列的值应是当前 DataFrame 在字典中的键（即来源名称）。pl.lit() 函数用于将一个 Python 值转换为 Polars 的字面量 Series，确保新列的所有行都包含相同的来源名称。
4. 将所有经过修改的 DataFrame 收集到一个列表中。
5. 使用 pl.concat 函数将这个列表中的所有 DataFrame 垂直合并。

示例数据：

假设我们有以下一个包含两个 Polars DataFrame 的字典：

NameGPT名称生成器

免费AI公司名称生成器，AI在线生成企业名称，注册公司名称起名大全。

0

查看详情

import polars as pl

dcty = {
    "df1": pl.DataFrame({'col1': [1, 2], 'col2': ["a", "b"]}),
    "df2": pl.DataFrame({'col1': [3, 4], 'col2': ["c", "d"]}),
}

print("原始字典中的 DataFrame:")
print("df1:")
print(dcty["df1"])
print("\ndf2:")
print(dcty["df2"])

解决方案代码：

# 使用列表推导式为每个 DataFrame 添加 'sheet' 列，然后进行合并
combined_df = pl.concat(
    [df.with_columns(sheet=pl.lit(name)) for name, df in dcty.items()]
)

print("\n合并后的 DataFrame:")
print(combined_df)

代码解析：

• for name, df in dcty.items(): 这部分遍历了字典 dcty 中的所有键值对。name 会依次是 "df1" 和 "df2"，df 则是对应的 Polars DataFrame。
• df.with_columns(sheet=pl.lit(name)): 这是解决方案的关键。
  • df.with_columns(...) 用于向 DataFrame 添加新列或修改现有列。
  • sheet=... 定义了新列的名称为 "sheet"。
  • pl.lit(name) 创建了一个字面量 Series。这意味着对于当前 df 中的所有行，新列 "sheet" 的值都将是当前的 name（例如 "df1" 或 "df2"）。
• [...]: 列表推导式将所有经过 with_columns 处理后的 DataFrame 收集到一个列表中。
• pl.concat(...): 最后，pl.concat 函数接收这个 DataFrame 列表，并默认进行垂直合并（how="vertical"）。

预期输出：

运行上述代码，将得到以下合并后的 DataFrame：

shape: (4, 3)
┌──────┬──────┬───────┐
│ col1 ┆ col2 ┆ sheet │
│ ---  ┆ ---  ┆ ---   │
│ i64  ┆ str  ┆ str   │
╞══════╪══════╪═══════╡
│ 1    ┆ a    ┆ df1   │
│ 2    ┆ b    ┆ df1   │
│ 3    ┆ c    ┆ df2   │
│ 4    ┆ d    ┆ df2   │
└──────┴──────┴───────┘

可以看到，col1 和 col2 列的数据已成功合并，并且新增了一个 sheet 列，准确地记录了每一行数据来自哪个原始 DataFrame。

注意事项

1. Schema 一致性： pl.concat 在垂直合并 DataFrame 时，要求所有参与合并的 DataFrame 具有兼容的 Schema（即相同的列名和数据类型）。如果 Schema 不一致，Polars 可能会尝试进行类型推断或抛出错误。在实际应用中，如果来源 DataFrame 的 Schema 可能不一致，您可能需要在合并前进行数据清洗、列选择或类型转换，以确保兼容性。
2. 性能考量： 对于包含大量 DataFrame 的字典，这种方法通常是高效的。Polars 在内部对这些操作进行了优化。
3. 列名冲突： 确保您为来源名称指定的新列名（例如本例中的 sheet）不会与原始 DataFrame 中已有的列名冲突，否则现有列可能会被覆盖。

总结

通过本教程，我们学习了如何利用 Polars 的强大功能，将一个字典中的多个 DataFrame 垂直合并为一个单一的 DataFrame，并在合并过程中智能地添加一个新列以标识每个数据行的原始来源。这种方法在处理来自多个数据源的数据时非常实用，能够帮助我们更好地管理和理解数据。掌握 with_columns 和 pl.lit 与 pl.concat 的结合使用，将极大地提升您在 Polars 中的数据处理效率和灵活性。

以上就是在 Polars 中从字典合并 DataFrame 并保留来源名称的详细内容，更多请关注php中文网其它相关文章！