优化Databricks中Python拆分字符串生成Pandas DataFrame的速度

在Databricks环境中，使用Python从大型PDF文件提取数据并转换为Pandas DataFrame时，字符串拆分及DataFrame拼接环节耗时约5小时，以下是针对性的优化方案：

一、PDF数据提取阶段优化 #

原始代码中字符串拼接、PDF读取方式存在性能瓶颈，优化点如下：

• 替换低效的字符串拼接方式：
  原始代码用extracted_string += ...拼接文本，由于Python字符串是不可变对象，每次拼接都会生成新字符串，内存开销大且速度慢。改用列表收集每页文本，最后一次性join，性能提升显著。

• 使用更高效的PDF读取库/API：
  PyPDF2的PdfFileReader已被标记为过时，建议使用新版的PdfReader；对于表格类PDF，pdfplumber的文本提取精度和速度更优（需先安装：pip install pdfplumber）。

• 简化循环逻辑：用for循环遍历页码，比while循环更简洁高效。

优化后的PDF提取代码：

import os
import re
# 可选：用pdfplumber替换PyPDF2
import pdfplumber
# 或者保留PyPDF2但用新版API
# from PyPDF2 import PdfReader

data_directory = "/path/to/your/pdf/folder"
all_table_lines = []

for pdf in os.listdir(data_directory):
    pdf_path = os.path.join(data_directory, pdf)
    # 方式1：用pdfplumber读取
    with pdfplumber.open(pdf_path) as pdf_obj:
        extracted_text = ""
        for page in pdf_obj.pages:
            extracted_text += page.extract_text() or ""
    
    # 方式2：用PyPDF2新版API
    # with open(pdf_path, 'rb') as file:
    #     reader = PdfReader(file)
    #     extracted_text = ""
    #     for page in reader.pages:
    #         extracted_text += page.extract_text() or ""

    # 匹配目标行
    regex_date = r"\d{2}/\d{2}/\d{4}[^\n]*"
    table_lines = re.findall(regex_date, extracted_text)
    all_table_lines.extend(table_lines)

二、DataFrame生成阶段优化 #

原始代码在循环中反复concat DataFrame是性能杀手，优化核心是先收集所有数据到列表，最后一次性生成DataFrame：

• 避免循环内拼接DataFrame：每次pd.concat都会创建新的DataFrame对象，循环次数越多，开销越大。改用列表存储所有拆分后的数据行，最后统一转换。
• 批量处理字符串拆分：用嵌套列表推导式一次性处理所有行的拆分，比循环逐行处理更快。
• 统一设置列名：无需在循环内重命名列，生成DataFrame时直接指定列名。

优化后的DataFrame生成代码：

import pandas as pd

# 批量拆分所有行的数据
all_rows = []
for line in all_table_lines:
    # 按换行拆分每行，再按空格拆分字段（注意：若字段含空格需调整拆分逻辑）
    rows_in_line = [x.split(' ') for x in line.split('\n') if x.strip()]  # 过滤空行
    all_rows.extend(rows_in_line)

# 一次性生成DataFrame并设置列名
table = pd.DataFrame(all_rows, columns=[
    'date_of_import', 'entry_num', 'warehouse_code_num', 
    'declarant_ref_num', 'declarant_EORI_num', 'VAT_due'
])

# 导出为CSV
table.to_csv("/path/to/output.csv", index=False)

额外优化建议 #

• 处理字段含空格的情况：如果原始数据中存在带空格的字段（如VAT_due可能包含金额符号和空格），直接用split(' ')会拆分错误，可改用正则表达式拆分，例如re.split(r'\s{2,}', x)（按连续空格拆分）。
• 利用Databricks分布式计算：如果PDF文件数量极多，可将任务拆分为Spark RDD或DataFrame的分布式任务，利用集群资源并行处理。
• 数据类型提前指定：生成DataFrame时通过dtype参数指定各列的数据类型，避免Pandas自动推断类型的开销。

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