本文介绍了如何将包含 CASE 表达式作为 JOIN 条件的 SQL 查询转换为 Pandas DataFrame 操作。核心思路是利用 Pandas 的 query() 方法预先过滤用于 JOIN 的 DataFrame，然后使用 merge() 方法进行连接。通过这种方式，可以有效地模拟 SQL 中 CASE WHEN 语句在 JOIN 条件中的行为，实现数据关联和筛选。
在 SQL 中，我们经常会遇到需要在 JOIN 条件中使用 CASE 表达式的情况，这允许我们根据不同的条件选择不同的连接键。当需要将这样的 SQL 查询转换为 Pandas 操作时，直接使用 pd.merge() 可能会遇到困难。以下介绍一种常用的方法，通过预先过滤 DataFrame 来模拟 CASE 表达式的行为。
问题描述
假设我们有以下 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)
这个查询的目的是将 table_a 和 table_b 进行左连接，连接条件是 country 相等，并且当 table_b 的 country 是 'Europe' 时，table_b 的 year 必须是 2022，否则 table_b 的 year 必须是 2023。
解决方案
由于 Pandas 的 merge() 方法不直接支持 CASE 表达式，我们可以首先使用 query() 方法对 table_b 进行过滤，然后再进行 merge() 操作。
- 创建示例 DataFrame
首先，我们创建两个示例 DataFrame，模拟 table_a 和 table_b：
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
接下来，我们使用 query() 方法对 table_b 进行过滤，模拟 CASE 表达式的行为：
table_b_filtered = table_b.query("(country == 'Europe' and year == 2022) or (country != 'Europe' and year == 2023)")这个 query() 表达式等价于 SQL 中的 CASE 表达式。它选择了 country 为 'Europe' 且 year 为 2022 的行，或者 country 不为 'Europe' 且 year 为 2023 的行。
- 使用 merge() 进行左连接
最后，我们使用 merge() 方法将 table_a 和过滤后的 table_b 进行左连接：
output = (
table_a.merge(
table_b_filtered,
on=['country'], how='left', suffixes=('', '_'))
[['year', 'country', 'amount']]
)
print(output)在 merge() 方法中，我们指定 on=['country'] 作为连接键，how='left' 表示进行左连接。suffixes=('', '_') 用于处理连接后相同列名的情况，这里我们只保留了 amount 列。最后，我们选择需要的列 [['year', 'country', 'amount']]。
完整代码示例
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]
})
output = (
table_a.merge(
table_b.query("(country == 'Europe' and year == 2022) or (country != 'Europe' and year == 2023)"),
on=['country'], how='left', suffixes=('', '_'))
[['country', 'year', 'amount']]
)
print(output)输出结果
country year amount 0 Europe 2022 50.0 1 Europe 2020 50.0 2 USA 2023 40.0 3 Africa 2021 NaN
总结
通过使用 query() 方法预先过滤 DataFrame，我们可以有效地模拟 SQL 中 CASE 表达式在 JOIN 条件中的行为。这种方法可以帮助我们将复杂的 SQL 查询转换为 Pandas 操作，从而更好地利用 Pandas 的数据处理能力。
注意事项
- query() 方法的性能可能受到数据量大小的影响。对于大型 DataFrame，可以考虑使用其他优化方法，例如使用 np.where() 或自定义函数。
- 在复杂的 CASE 表达式中，query() 表达式可能会变得难以阅读和维护。建议将复杂的逻辑拆分成多个简单的步骤，以提高代码的可读性。
- 确保 query() 表达式的逻辑与 SQL 中的 CASE 表达式完全一致，以避免出现错误的结果。
