本文介绍了如何使用 Pandas 实现 SQL 中带有 CASE 表达式的 JOIN 操作。通过先对 DataFrame 进行过滤，再进行合并，可以有效地模拟 SQL 中 CASE WHEN 的逻辑，从而实现复杂条件下的数据连接。本文将提供详细的代码示例和步骤，帮助读者理解并应用该方法。
在 SQL 中，CASE 表达式常用于 JOIN 条件中，以实现更灵活的连接逻辑。然而，Pandas 的 pd.merge() 函数直接实现这种复杂的 JOIN 条件并不直观。本文将介绍一种通过预先过滤 DataFrame 的方式，来模拟 SQL 中 CASE 表达式在 JOIN 中的应用。
问题描述
假设我们有两个 DataFrame，table_a 和 table_b，需要根据以下 SQL 查询进行连接：
SELECT a.year, a.country, b.amount FROM table_a a LEFT JOIN table_b b ON a.country=b.country AND (CASE WHEN b.country = 'Europe' THEN b.year = 2022 ELSE b.year = 2023 END)
解决方案
由于 CASE 表达式用于在 JOIN 之前过滤 table_b，我们可以先使用 Pandas 的 query() 方法过滤 table_b，然后再进行 merge() 操作。
示例代码
首先，创建示例 DataFrame：
import pandas as pd
table_a = pd.DataFrame({
'country': ['Europe', 'Europe', 'USA', 'Africa'],
'year': [2022, 2020, 2023, 2021]
})
table_b = pd.DataFrame({
'country': ['Europe', 'USA', 'Africa', 'USA', 'Europe'],
'year': [2023, 2022, 2022, 2023, 2022],
'amount': [10, 20, 30, 40, 50]
})接下来，使用 query() 方法过滤 table_b：
table_b_filtered = table_b.query("(country == 'Europe' and year == 2022) or (country != 'Europe' and year == 2023)")然后，使用 merge() 方法进行左连接：
output = (
table_a.merge(
table_b_filtered,
on=['country'], how='left', suffixes=('', '_'))
[['year', 'country', 'amount']]
)
print(output)代码解释
- table_b.query("(country == 'Europe' and year == 2022) or (country != 'Europe' and year == 2023)")：这行代码使用 query() 方法根据 CASE 表达式的逻辑过滤 table_b。query() 方法允许我们使用字符串表达式来指定过滤条件。
- table_a.merge(table_b_filtered, on=['country'], how='left', suffixes=('', '_'))：这行代码使用 merge() 方法将 table_a 和过滤后的 table_b_filtered 进行左连接。on=['country'] 指定连接键为 country 列。how='left' 表示左连接。suffixes=('', '_') 用于处理连接后可能出现的重复列名。
- [['year', 'country', 'amount']]：这行代码选择最终输出的列。
输出结果
运行上述代码，将得到以下结果：
year country amount 0 2022 Europe 50.0 1 2020 Europe 50.0 2 2023 USA 40.0 3 2021 Africa NaN
注意事项
- query() 方法的性能可能不如直接使用布尔索引。如果 DataFrame 非常大，可以考虑使用布尔索引来提高性能。
- 在复杂的 CASE 表达式中，可以使用多个 query() 方法进行分步过滤，以提高代码的可读性。
- 确保 merge() 方法的连接键选择正确，避免出现错误的结果。
总结
通过预先过滤 DataFrame，可以有效地模拟 SQL 中 CASE 表达式在 JOIN 中的应用。这种方法简单易懂，适用于大多数情况。在实际应用中，需要根据具体情况选择合适的过滤方法，并注意性能优化。
