Pandas DataFrame中JSON字符串反向查找的优化方案咨询

原始DataFrame结构 #

给定不可修改的Pandas DataFrame，结构如下：

import pandas as pd

json_str_01 = '''[{"my_key": 1, "my_values": ["3_13"]}, {"my_key": 2, "my_values": ["3_18", "2_12"]}, {"my_key": 3, "my_values": ["2_10", "2_17"]}]'''
json_str_02 = '''[{"my_key": 1, "my_values": ["2_17"]}, {"my_key": 2, "my_values": ["2_12", "3_16"]}, {"my_key": 3, "my_values": ["3_13", "2_8"]}, {"my_key": 4, "my_values": ["3_18", "3_17"]}]'''

data = {
  "search_str": ["3_13", "2_8"],
  "json_data_str": [json_str_01, json_str_02]
}

df = pd.DataFrame(data)

每行包含两列：

• search_str：待查找的目标字符串
• json_data_str：可解析为JSON数组的字符串，数组中每个对象包含my_key（整数）和my_values（字符串列表）

需求 #

新增一列result，返回search_str在对应行my_values列表中首次出现时对应的my_key；若无匹配返回None（题目说明至少存在一个匹配）。

现有可行代码 #

以下代码可得到正确结果，但依赖逐行循环，大数据量下性能较差：

def json_reverse_lookup(search_str, json_data_str):
    df_json_data = pd.read_json(json_data_str)

    for idx, row in df_json_data.iterrows():
        if search_str in row[1]:
            return row[0]
    
    return None

df['result'] = df.apply(lambda x: json_reverse_lookup(x['search_str'], x['json_data_str']), axis=1)
print(df)

优化方案 #

方案一：结构化展开+分组匹配 #

通过解析JSON并展开成结构化数据，利用Pandas的分组匹配快速定位首个匹配键：

import pandas as pd
import json

# 解析JSON并展开，保留原行索引
df_json_expanded = df['json_data_str'].apply(lambda x: pd.json_normalize(json.loads(x))).explode('my_values').reset_index()

# 合并原数据的search_str，关联匹配关系
df_merged = df_json_expanded.merge(df[['search_str']], left_on='index', right_index=True)

# 按原行索引分组，取首个匹配的my_key
df_result = df_merged[df_merged['my_values'] == df_merged['search_str']].groupby('index')['my_key'].first().reset_index()

# 关联回原DataFrame
df = df.merge(df_result, left_index=True, right_on='index', how='left').drop('index', axis=1).rename(columns={'my_key': 'result'})

print(df)

方案二：字典映射（高效推荐） #

预先为每行JSON构建值→首个键的映射字典，再通过哈希查询直接获取结果，避免循环：

import pandas as pd
import json

# 为单个JSON字符串构建值到首个my_key的映射
def build_value_key_map(json_str):
    items = json.loads(json_str)
    value_map = {}
    for item in items:
        current_key = item['my_key']
        for val in item['my_values']:
            if val not in value_map:  # 仅保留首次出现的键
                value_map[val] = current_key
    return value_map

# 生成每行的映射字典
df['value_map'] = df['json_data_str'].apply(build_value_key_map)

# 直接通过映射获取结果
df['result'] = df.apply(lambda x: x['value_map'].get(x['search_str']), axis=1)

# 清理中间列
df = df.drop('value_map', axis=1)

print(df)

优化说明 #

• 两种方案均避免了iterrows和逐行循环的低效操作，利用Pandas向量化特性或哈希表查询提升性能。
• 方案二在数据量较大时优势更明显：每个JSON仅解析一次，且查找操作是O(1)的哈希查询，整体时间复杂度更低。

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