在Pandas/Dask中实现ID与对应国家的匹配关联

问题背景 #

给定如下格式的Pandas数据集：

import pandas as pd
# initialize list of lists
data = [[2, [4], "['RU']"], [4, [9,18,24], "['US' 'CN']"], [6, [],  "['']"],[7, [2],  "['DE' 'IT']"],[9, [4],  "['US']"],[14, [18,6],  "['SA']"],[18, [17],  "['GB' 'KR']"]]
  
# Create the pandas DataFrame
df = pd.DataFrame(data, columns=['docdb', 'cited_docdb','country'])

其中country列存储docdb对应ID的国家信息，cited_docdb中的部分ID未出现在docdb列表中。需要新增desired列，包含cited_docdb中所有存在于docdb的ID对应的国家信息，且代码需适配大数据集。期望结果如下：

#initialize list of lists
data_result = [[2, [4], "['RU']", "['US' 'CN']"], [4, [9,18,24], "['US' 'CN']","['US','GB' 'KR','']"], [6, [],  "['']", "['']"],[7, [2],  "['DE' 'IT']", "['']"],[9, [4],  "['US']", "['US' 'CN']"],[14, [18,6],  "['SA']", "['GB' 'KR', '']"],[18, [17],  "['GB' 'KR']","['']"]]
  
# Create the pandas DataFrame
df_result = pd.DataFrame(data_result, columns=['docdb', 'cited_docdb','country', 'desired'])

高效实现方案 #

针对大数据集，优先使用Pandas的矢量化操作（避免逐行迭代），核心思路是利用映射+分组聚合：

步骤1：创建DocID到国家信息的映射 #

先建立docdb与country的一对一映射，用于快速查询引用文档的国家信息：

# 创建docdb到country的映射Series
docdb_to_country = df.set_index('docdb')['country']

步骤2：拆分引用文档列表并匹配国家信息 #

将cited_docdb列的列表拆分为多行，然后通过映射匹配对应国家，不存在的ID填充默认空值：

# 拆分cited_docdb列表为多行
exploded_df = df.explode('cited_docdb', ignore_index=True)

# 匹配国家信息，不存在的ID填充"['']"
exploded_df['matched_country'] = exploded_df['cited_docdb'].map(docdb_to_country).fillna("['']")

步骤3：分组聚合生成desired列 #

按docdb分组，将每个分组内的匹配国家信息合并为目标格式的字符串：

def format_desired(country_strings):
    # 处理每个国家字符串，提取内部内容并格式化
    formatted_parts = []
    for s in country_strings:
        inner_content = s.strip('[]')
        if not inner_content:
            formatted_parts.append("''")
        else:
            # 保留原始的空格分隔格式（如'US' 'CN'）
            formatted_parts.append(f"'{inner_content}'")
    # 拼接成最终的字符串格式
    return f"[{', '.join(formatted_parts)}]"

# 分组聚合生成desired列
desired_series = exploded_df.groupby('docdb')['matched_country'].agg(format_desired)

步骤4：合并回原数据集 #

将生成的desired列合并到原始DataFrame中：

# 合并数据，确保所有原始行都保留
df = df.merge(desired_series.rename('desired'), on='docdb', how='left')

完整代码 #

import pandas as pd

# 初始化原始数据
data = [[2, [4], "['RU']"], [4, [9,18,24], "['US' 'CN']"], [6, [],  "['']"],[7, [2],  "['DE' 'IT']"],[9, [4],  "['US']"],[14, [18,6],  "['SA']"],[18, [17],  "['GB' 'KR']"]]
df = pd.DataFrame(data, columns=['docdb', 'cited_docdb','country'])

# 步骤1：创建映射
docdb_to_country = df.set_index('docdb')['country']

# 步骤2：拆分并匹配
exploded_df = df.explode('cited_docdb', ignore_index=True)
exploded_df['matched_country'] = exploded_df['cited_docdb'].map(docdb_to_country).fillna("['']")

# 步骤3：分组格式化
def format_desired(country_strings):
    formatted_parts = []
    for s in country_strings:
        inner_content = s.strip('[]')
        formatted_parts.append("''" if not inner_content else f"'{inner_content}'")
    return f"[{', '.join(formatted_parts)}]"

desired_series = exploded_df.groupby('docdb')['matched_country'].agg(format_desired)

# 步骤4：合并
df = df.merge(desired_series.rename('desired'), on='docdb', how='left')

# 查看结果
print(df)

方案优势 #

• 高效性：全程使用Pandas的矢量化操作和分组聚合，避免了iterrows()这类低效的逐行迭代，处理百万级数据集性能更优。
• 可扩展性：映射和分组逻辑可以轻松适配更大规模的数据集，无需修改核心逻辑。
• 格式一致性：严格匹配需求中的字符串格式，包括空值处理和原始国家信息的格式保留。

内容的提问来源于stack exchange，提问作者Lusian