如何在Pandas中分块处理数据？

Question

我有一个函数：

def extract_named_entities(note):
    """
    Extract the named entities identified in a given note.
    """
    doc = nlp(note)
    return [ent.text for ent in doc.ents]

df['named_entities'] = df['NOTE'].apply(extract_named_entities)

e_df = df.explode('named_entities').reset_index(drop=True)

在这个函数中，df['NOTE']的每一行都包含一个2500字的段落。我想优化这个函数，因为它在处理5到10行数据时运行得很快。但是，我有6400行数据，这样就花了很多时间。

有没有办法使用分块处理或者其他优化技巧，或者能不能避免使用列表呢？

Answer 1

我假设你代码里的 nlp 实际上是一个 SpaCy 模型。如果是这样的话，你的函数并不是在提取实体，而是在提取一个“NOTE”的 SpaCy 词元。如果每段文本大约有 2500 个单词，那么每一行就会有超过 2500 个词元，并且每一行都会调用 nlp(...)（这个过程比较慢）。

当你把这些列表展开时，你会得到一个包含 6400 * 2500 行的数据框（如果你保留其他列的话，这样会得到一个非常庞大且重复的数据框）。

pandas.apply 的调用只是遍历每一行——这就像一个循环。要让提取函数（保持不变）更快，唯一的方法就是并行处理。简单地将函数向量化（也就是让它对整列进行操作，而不是单独对每个单元格操作）是不够的，因为真正的工作是在 nlp(...) 调用内部完成的，而你无法改变 nlp 的实现。

如果你只想要词元，那么可以通过将提取函数定义为以下方式来优化它：

def extract_tokens(text):
    return [x.text for x in nlp.tokenize(text)]

在大约 6800 行，每行最多约 3200 个单词（平均约 730 个单词）的情况下，这个过程总共运行不到 6 秒（在 Apple M2 笔记本上），而 nlp(text) 的速度明显慢得多（慢了超过 10 倍）。