本教程详细介绍了如何在polars中高效地将存储在字典中的多个dataframe合并成一个单一的dataframe，同时为每行数据添加一个新列，以标识其原始来源（即字典的键）。通过结合使用列表推导式、`with_columns`、`pl.lit`和`pl.concat`函数，用户可以优雅地实现这一常见的数据整理任务，确保数据完整性和可追溯性。
Polars中合并字典DataFrame并保留来源信息
在数据处理过程中，我们经常会遇到需要从多个来源（例如不同的Excel工作表、数据库查询结果）导入数据，并将它们存储在一个Python字典中，其中字典的键代表来源名称，值是对应的Polars DataFrame。此时，一个常见的需求是将这些独立的DataFrame按行合并成一个统一的DataFrame，并且在新合并的表中保留一个列来标识每行数据最初属于哪个DataFrame。本文将深入探讨如何在Polars中优雅且高效地实现这一目标。
问题场景与挑战
假设我们有一个字典dcty，其中包含了多个Polars DataFrame，每个DataFrame都代表一个独立的数据集：
import polars as pl
dcty = {
"df1": pl.DataFrame({'col1': [1, 2], 'col2': ["a", "b"]}),
"df2": pl.DataFrame({'col1': [3, 4], 'col2': ["c", "d"]}),
}我们期望的输出是一个包含所有数据，并且多了一个名为sheet的列，用于指示数据来源的DataFrame：
shape: (4, 3) ┌──────┬──────┬──────┐ │ col1 ┆ col2 ┆ sheet│ │ --- ┆ --- ┆ --- │ │ i64 ┆ str ┆ str │ ╞══════╪══════╪══════╡ │ 1 ┆ a ┆ df1 │ │ 2 ┆ b ┆ df1 │ │ 3 ┆ c ┆ df2 │ │ 4 ┆ d ┆ df2 │ └──────┴──────┴──────┘
直接使用pl.concat函数可以垂直合并DataFrame列表，但它无法自动添加来源标识列。如果我们将字典转换为列表，会丢失原始的DataFrame名称信息。因此，我们需要一种方法在合并之前，为每个DataFrame添加一个包含其名称的新列。
核心解决方案
Polars提供了一个简洁而强大的解决方案，它结合了列表推导式、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)运行上述代码将产生预期的输出：
shape: (4, 3) ┌──────┬──────┬───────┐ │ col1 ┆ col2 ┆ sheet │ │ --- ┆ --- ┆ --- │ │ i64 ┆ str ┆ str │ ╞══════╪══════╪═══════╡ │ 1 ┆ a ┆ df1 │ │ 2 ┆ b ┆ df1 │ │ 3 ┆ c ┆ df2 │ │ 4 ┆ d ┆ df2 │ └──────┴──────┴───────┘
解决方案详解
这个解决方案的核心在于通过列表推导式对字典中的每个DataFrame进行预处理，然后将处理后的DataFrame列表传递给pl.concat。
dcty.items(): 这个方法用于遍历字典dcty中的所有键值对。在每次迭代中，name将是字典的键（例如"df1", "df2"），而df将是对应的Polars DataFrame。
-
df.with_columns(sheet=pl.lit(name)): 这是关键的一步，它为每个DataFrame动态地添加一个新列。
- with_columns(): Polars DataFrame的一个方法，用于添加、修改或删除列。它接受一个或多个表达式作为参数。
- sheet=...: 指定新列的名称为sheet。
- pl.lit(name): pl.lit()函数用于创建一个字面量表达式。在这里，它将当前的name（即字典的键，如"df1"）作为一个常量值，应用于新列sheet的所有行。这意味着，来自"df1"的所有行都会在sheet列中得到"df1"这个值。
[...] (列表推导式): 列表推导式[df.with_columns(sheet=pl.lit(name)) for name, df in dcty.items()]会生成一个新的DataFrame列表。这个列表中的每个DataFrame都已经是经过修改的，即它们都包含了一个名为sheet的新列，其中填充了其原始的字典键。
pl.concat(...): 最后，pl.concat()函数接收这个包含所有已修改DataFrame的列表，并将它们垂直（按行）合并成一个单一的DataFrame。pl.concat默认执行垂直合并，因此不需要显式指定how="vertical"。
通过这种方式，我们不仅成功地合并了所有DataFrame，而且巧妙地利用了Polars的表达式系统，在合并过程中保留了宝贵的来源信息。
注意事项与最佳实践
- Schema一致性: 为了确保pl.concat顺利执行并得到预期的结果，所有待合并的DataFrame应该具有兼容的列结构和数据类型。如果列名不一致，pl.concat会尝试进行对齐，可能会导致额外的null值。如果数据类型不兼容，Polars可能会尝试隐式类型转换，但也可能抛出错误。在合并前检查并标准化DataFrame的schema是一个好习惯。
- 性能: Polars的pl.concat函数是高度优化的，对于大量DataFrame或大型DataFrame的合并操作，其性能通常非常出色。列表推导式在Python中也是一种高效的构建列表的方式。
- 内存管理: 尽管Polars在内存效率方面表现优秀，但合并大量大型DataFrame仍然可能消耗大量内存。确保系统有足够的RAM来处理合并操作。
-
错误处理: 在实际应用中，如果字典中的DataFrame可能为空或格式不一致，建议在列表推导式中加入额外的逻辑进行过滤或处理，例如：
# 过滤掉空的DataFrame combined_df = pl.concat([ df.with_columns(sheet=pl.lit(name)) for name, df in dcty.items() if not df.is_empty() ])
总结
将Polars字典中的多个DataFrame合并成一个带有来源标识的单一DataFrame，是一个常见的但有时令人困惑的数据处理任务。通过本教程介绍的方法，即利用列表推导式结合with_columns和pl.lit来预处理每个DataFrame，再通过pl.concat进行最终合并，我们能够以一种高效、清晰且富有表达力的方式解决这一问题。掌握这一技巧，将大大提升您在Polars中处理复杂数据整合任务的能力。
