Polars 数据框字典合并与源信息保留教程

本教程详细介绍了如何在 polars 中高效地将存储在字典中的多个数据框按行合并为一个单一的数据框，并在此过程中自动添加一列以记录每个观测值来源于哪个原始数据框（即字典的键名）。通过结合列表推导式、`with_columns` 和 `pl.concat` 函数，我们能够优雅地解决这一常见的数据整合需求，同时保持代码的简洁性和执行效率。

在数据处理工作中，我们经常会遇到需要从多个来源（例如不同的 Excel 工作表、CSV 文件或数据库查询结果）导入数据，并将它们存储在一个字典中，其中字典的键通常代表数据的来源名称。随后，一个常见的需求是将这些独立的数据框按行合并成一个大的数据框，同时保留原始数据框的名称作为新数据框中的一列，以便追踪数据的来源。

Polars 作为一种高性能的数据框库，提供了强大的工具来处理这类任务。虽然它没有像 R Tidyverse 中 bind_rows(.id = "ID") 那样直接的单函数解决方案，但通过组合其核心功能，我们可以实现同样甚至更灵活的效果。

问题场景与目标

假设我们有一个 Polars 数据框的字典，其中每个键是数据框的名称，值是对应的 Polars 数据框。我们的目标是：

1. 将字典中的所有数据框按行垂直堆叠。
2. 在合并后的数据框中添加一个新列（例如 sheet），该列的值是每个观测值所来自的原始数据框的名称（即字典的键）。

以下是一个示例数据字典：

import polars as pl

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

我们期望的输出结果是一个包含所有行，并额外带有 sheet 列的单一数据框：

怪兽AI数字人

数字人短视频创作，数字人直播，实时驱动数字人

44

查看详情

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

解决方案详解

Polars 提供了 pl.concat 函数用于合并数据框。默认情况下，pl.concat 执行垂直合并（how="vertical"），这正是我们需要的。然而，pl.concat 接收的是一个数据框列表，直接将字典的值转换为列表会丢失原始数据框的名称信息。因此，关键在于在合并之前，为每个数据框添加一个包含其名称的新列。

我们可以通过结合列表推导式、DataFrame.with_columns() 方法和 pl.lit() 函数来实现这一点。

核心步骤：

1. 遍历字典： 使用列表推导式遍历字典中的每个键值对（name, df）。
2. 添加源名称列： 对于字典中的每一个数据框 df，使用 df.with_columns(sheet=pl.lit(name)) 添加一个名为 sheet 的新列。
   • with_columns()：这是 Polars 中用于添加或修改列的强大方法。
   • pl.lit(name)：pl.lit() 函数用于创建一个字面量表达式。在这里，它会为新添加的 sheet 列的每一行填充当前数据框的名称 name。
3. 执行垂直合并： 将经过上述处理的所有数据框组成的列表传递给 pl.concat() 函数，完成最终的合并。

示例代码：

import polars as pl

# 示例数据字典
dcty = {
    "df1": pl.DataFrame({'col1': [1, 2], 'col2': ["a", "b"]}),
    "df2": pl.DataFrame({'col1': [3, 4], 'col2': ["c", "d"]}),
}

# 解决方案
combined_df = pl.concat([
    df.with_columns(sheet=pl.lit(name)) 
    for name, df in dcty.items()
])

# 打印结果
print(combined_df)

运行上述代码将产生预期的输出：

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

注意事项与最佳实践

• 列名一致性： 进行垂直合并时，建议所有待合并的数据框具有相同的列名和兼容的数据类型。如果列名不一致，pl.concat 会自动填充 null 值以对齐列。如果数据类型不兼容，Polars 会尝试进行类型推断或强制转换，这可能导致意外的结果或错误。
• 性能： 这种基于列表推导式和 pl.concat 的方法在 Polars 中是高效且惯用的。Polars 内部会优化这些操作，尤其是在处理大量数据时，其性能通常优于传统的循环追加方法。
• 灵活性： with_columns 方法非常灵活，你不仅可以添加源名称，还可以在合并前对每个数据框执行其他预处理操作，例如重命名列、过滤行等。
• pl.lit() 的作用： pl.lit() 是一个关键函数，它将一个 Python 值转换为 Polars 表达式，确保该值在整个新列中被广播（重复）到每一行。

总结

本教程展示了在 Polars 中如何优雅地将一个数据框字典合并为一个单一的数据框，并在此过程中保留原始数据框的名称作为新列。通过利用列表推导式、DataFrame.with_columns() 和 pl.lit() 函数，我们能够构建出既简洁又高效的代码，完美解决了在数据整合场景中追踪数据来源的常见需求。掌握这种模式对于 Polars 用户来说至关重要，它能够帮助你更有效地组织和分析复杂数据集。

以上就是Polars 数据框字典合并与源信息保留教程的详细内容，更多请关注php中文网其它相关文章！