本文详细介绍了如何在 polars 中高效合并存储在字典中的多个数据框。通过结合列表推导式、`with_columns` 方法和 `pl.concat` 函数，可以轻松地将这些数据框按行堆叠成一个单一的数据框，同时自动添加一个新列以标识每行数据最初来源于哪个数据框，从而保留重要的元数据信息，提升数据处理的灵活性和可追溯性。
Polars 中合并字典数据框并保留来源信息
在数据分析和处理中，我们经常会遇到需要从多个来源（例如不同的 Excel 工作表、CSV 文件或 API 响应）导入数据，并将它们存储在一个字典中，其中字典的键代表数据的来源名称，值则是对应的 Polars 数据框。一个常见的需求是将这些分散的数据框合并成一个统一的 Polars 数据框，同时在合并后的数据中增加一个列，用于标识每行数据最初来自哪个数据框。
虽然 Polars 提供了强大的 pl.concat 函数用于数据框的垂直或水平合并，但它本身并不直接支持在合并过程中自动添加一个标识来源的列。如果简单地将字典中的数据框转换为列表再进行合并，我们将丢失原始的键（即来源名称）。本教程将展示如何优雅地解决这一问题，实现类似于 R Tidyverse 中 bind_rows(.id = "ID") 的功能。
核心问题与挑战
假设我们有一个包含多个 Polars 数据框的字典，如下所示：
import polars as pl
dcty = {
"df1": pl.DataFrame({'col1': [1, 2], 'col2': ["a", "b"]}),
"df2": pl.DataFrame({'col1': [3, 4], 'col2': ["c", "d"]}),
}我们的目标是将其合并成一个数据框，并新增一个名为 sheet 的列，其值对应于原始字典的键（"df1" 或 "df2"），预期输出如下：
shape: (4, 3) ┌──────┬──────┬───────┐ │ col1 ┆ col2 ┆ sheet │ │ --- ┆ --- ┆ --- │ │ i64 ┆ str ┆ str │ ╞══════╪══════╪═══════╡ │ 1 ┆ a ┆ df1 │ │ 2 ┆ b ┆ df1 │ │ 3 ┆ c ┆ df2 │ │ 4 ┆ d ┆ df2 │ └──────┴──────┴───────┘
直接使用 pl.concat(list(dcty.values())) 会合并数据框，但会丢失来源信息。
解决方案：结合 with_columns 和 pl.concat
解决此问题的关键在于在合并之前，为每个数据框添加一个包含其来源名称的新列。这可以通过列表推导式结合 pl.DataFrame.with_columns 方法和 pl.lit 函数来实现，然后将处理后的数据框列表传递给 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)代码解释：
- for name, df in dcty.items(): 这部分是一个列表推导式，它遍历字典 dcty 中的每一个键值对。name 将依次是字典的键（例如 "df1", "df2"），而 df 则是对应的值（即 Polars 数据框）。
-
df.with_columns(sheet=pl.lit(name)): 这是解决方案的核心。
- df.with_columns(): 这是 Polars 数据框的一个方法，用于添加或替换列。
- sheet=...: 指定要添加的新列的名称为 sheet。
- pl.lit(name): pl.lit() 函数用于创建一个字面量表达式。在这里，它将当前的 name（例如 "df1"）作为一个常量值，为当前数据框的所有行创建一个名为 sheet 的新列，并填充这个常量值。 通过这一步，每个数据框在被合并之前，都拥有了一个记录其来源的 sheet 列。
- pl.concat([...]): 最后，pl.concat 函数接收一个 Polars 数据框的列表，并默认进行垂直堆叠（how="vertical"）。由于列表中的每个数据框都已包含了 sheet 列，因此合并后的数据框将自然地包含这一列，并正确地标识每行数据的来源。
注意事项与最佳实践
- 模式一致性： 尽管 Polars 在合并时具有一定的类型推断和兼容性处理能力，但为了避免潜在的类型不匹配错误或意外的数据转换，建议在合并前确保所有数据框的列名和对应的数据类型尽可能保持一致。如果存在不一致，Polars 可能会尝试进行类型强制转换，或者在严格模式下抛出错误。
- 列顺序： pl.concat 会根据第一个数据框的列顺序来排列合并后的列。如果后续数据框的列顺序不同，Polars 会自动调整以匹配，缺失的列会用 null 填充。
- 性能： 对于包含大量数据框的字典，这种列表推导式结合 pl.concat 的方法通常是高效的，因为它利用了 Polars 底层的并行化和优化。
- 替代方案（不推荐此场景）： 如果不需要保留来源信息，可以直接 pl.concat(list(dcty.values()))。但当来源信息至关重要时，本文介绍的方法是首选。
总结
通过巧妙地结合 Polars 的 with_columns 方法和 pl.concat 函数，我们可以轻松地实现将字典中存储的多个数据框高效地合并成一个单一数据框，并自动添加一个标识原始来源的新列。这种方法不仅代码简洁、易于理解，而且在处理复杂数据集成任务时，极大地提升了数据的可追溯性和分析效率。掌握这一技巧，将使您在 Polars 数据处理中更加游刃有余。
