Azure Databricks中CSV转格式化Excel函数的优化咨询

我在Azure Databricks笔记本上运行Spark集群，写了一个读取CSV并生成带格式Excel工作簿的函数，但遇到驱动节点内存过载问题——Pandas和xlsxwriter导致RAM占用飙升。我试过三种优化思路但效果不好：

• 用Spark完全替代Pandas处理数据：但必须用xlsxwriter实现Excel格式化，这条路走不通
• 更换xlsxwriter的替代库：无法复刻现有格式化功能
• 并行处理DataFrame并避免存入列表：代码可读性下降，且没法合理利用Spark集群资源

附上当前函数代码，求优化建议：

!pip install xlsxwriter --quiet --disable-pip-version-check

import pandas as pd
import numpy as np
import xlsxwriter
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, round, to_date
from pyspark.sql.types import DecimalType, DateType, TimestampType
import shutil

def check_lists_length(*lists): 
  return not(all(len(lst) == len(lists[0]) for lst in lists)) 

def client_file_creator_from_csv(sheet_names, table_names, file_names, csv_list, output_path): 
  """
  Reads a list of CSV files and returns formatted Excel Workbooks in Azure Databricks
  
  Parameters:
  sheet_names (list): A list of strings representing worksheet names of the final workbooks
  table_names (list): A list of strings representing table headers of the final workbooks
  file_names (list): A list of strings representing file names for the final workbooks
  csv_list (list): A list of strings representing the input paths of the CSV files
  output_path (string) : A string representing the output path for the Excel workbooks
  """
  if check_lists_length(sheet_names, table_names, file_names, csv_list):
    raise Exception("The number of CSV files and file names and sheet names and table names must be the same")

  #Convert CSV file paths to dbfs format for use with dbutils and pyspark
  csv_list = ["dbfs:" + str.split(x,"dbfs")[1] for x in csv_list]
  #Convert output file paths to dbfs format for use with dbutils and pyspark
  path_api = "dbfs:" + str.split(output_path,"dbfs")[1]
  #Create output path in case it doesn't exist
  dbutils.fs.mkdirs(path_api)
  #Create temporary path in dbfs to store ExcelWriter objects
  temp_path = "/tmp/" + str.split(dbutils.notebook.entry_point.getDbutils().notebook().getContext().userName().get(), sep = '@')[0]
  #Create lists to store pandas dataframes and ExcelWriter objects
  dfs = list()
  writers = list()
  for i,csv in enumerate(csv_list):
    df = spark.read.option("header", True).option("nullValue", "null").option("delimiter", ",").option("quote", "\"").option("escape", "\"").option("multiLine", True).option("inferSchema", True).option("enforceSchema", True).option("mode", "DROPMALFORMED").csv(csv)
    if df.count() == 0:
      raise Exception('File does not have any rows') 
    dcml_cols = [f.name for f in df.schema.fields if isinstance(f.dataType, DecimalType)]
    time_cols = [f.name for f in df.schema.fields if isinstance(f.dataType, TimestampType)]
    for col_name in dcml_cols:
      df = df.withColumn(col_name, round(col(col_name).cast('double'),2))
    for col_name in time_cols:
      df = df.withColumn(col_name, to_date(col(col_name)))
    dfs.append(df.select("*").toPandas())
    writers.append(pd.ExcelWriter(f"{temp_path}_{i}.xlsx", engine="xlsxwriter"))

  for df, writer, sheet_name, table_name, file_name in zip(dfs,writers, sheet_names, table_names, file_names):
    #Since Excel has a row limit of 1,048,576 rows per worksheet, larger dataframes must be split across multiple worksheets
    split_df = np.array_split(df,(df.shape[0]//1000000)+1)
    workbook = writer.book
    #Add standard formatting as required
    top_border = workbook.add_format({'top': 2})
    bottom_border = workbook.add_format({'bottom': 2})
    left_border = workbook.add_format({'left': 2})
    right_border = workbook.add_format({'right': 2})
    top_left_border = workbook.add_format({'top': 2, 'left': 2})
    top_right_border = workbook.add_format({'top': 2, 'right': 2})
    bottom_left_border = workbook.add_format({'bottom': 2, 'left': 2})
    bottom_right_border = workbook.add_format({'bottom': 2, 'right': 2})
    table_title_format = workbook.add_format({'bold': True, 'font_color': 'white', 'bg_color': '#305496', 'align': 'left', 'valign': 'vcenter'})
    title_format = workbook.add_format({'bold': True, 'font_color': 'white', 'bg_color': 'black', 'align': 'left', 'valign': 'vcenter'})
    for j, df_part in enumerate(split_df):
      df_part.to_excel(writer, sheet_name = sheet_name + (("_"+ str(j)) if j > 0 else ""), index = False, header = False, startrow = 4, startcol = 1)
      worksheet = writer.sheets[sheet_name + (("_"+ str(j)) if j > 0 else "")]
      for col_num, value in enumerate(df_part.columns.values):
        worksheet.write(3,col_num+1,value,table_title_format)
      num_rows = df_part.shape[0]
      num_cols = df_part.shape[1]
      worksheet.conditional_format(3,2,3,num_cols-1,{"type": "no_errors", "format": top_border})
      worksheet.conditional_format(4,1,num_rows+2,1,{"type": "no_errors", "format": left_border})
      worksheet.conditional_format(4,num_cols,num_rows+2,num_cols,{"type": "no_errors", "format": right_border})
      worksheet.conditional_format(num_rows+3,2,num_rows+3,num_cols-1,{"type": "no_errors", "format": bottom_border})
      worksheet.conditional_format(3,1,3,1,{"type": "no_errors", "format": top_left_border})
      worksheet.conditional_format(3,num_cols,3,num_cols,{"type": "no_errors", "format": top_right_border})
      worksheet.conditional_format(num_rows+3,1,num_rows+3,1,{"type": "no_errors", "format": bottom_left_border})
      worksheet.conditional_format(num_rows+3,num_cols,num_rows+3,num_cols,{"type": "no_errors", "format": bottom_right_border})
      worksheet.hide_gridlines(2)
      worksheet.autofilter(3,1,num_rows+3,num_cols)
      worksheet.autofit()
      worksheet.set_column_pixels(0, 0, 20)
      worksheet.merge_range(1,1,1,num_cols, table_name + (("_"+ str(j)) if j > 0 else ""), title_format)
    workbook.close()
    shutil.copy(writer,output_path + file_name + ".xlsx")
    print(f"File {file_name}.xlsx created successfully")

优化方案及修改后代码 #

核心优化措施 #

1. 避免批量加载所有数据到内存：不再将所有CSV对应的Pandas DF存入列表，改为逐个处理单个CSV，处理完成后立即清理资源，降低内存峰值
2. 分批次读取Spark数据：用limit+offset分批次读取Spark数据，每次只加载当前工作表所需数据，避免一次性转换全量Pandas DF
3. 启用xlsxwriter常量内存模式：让xlsxwriter逐行写入数据，不缓存整个工作表到内存
4. 即时资源清理：处理完批次数据后删除对象并触发垃圾回收，处理完文件后删除临时文件

修改后的函数代码 #

!pip install xlsxwriter --quiet --disable-pip-version-check

import pandas as pd
import numpy as np
import xlsxwriter
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, round, to_date
from pyspark.sql.types import DecimalType, DateType, TimestampType
import shutil
import gc
import os

def check_lists_length(*lists): 
  return not(all(len(lst) == len(lists[0]) for lst in lists)) 

def client_file_creator_from_csv(sheet_names, table_names, file_names, csv_list, output_path): 
  """
  读取CSV文件并在Azure Databricks中生成格式化Excel工作簿
  
  参数:
  sheet_names (list): 最终工作簿的工作表名称列表
  table_names (list): 最终工作簿的表头名称列表
  file_names (list): 最终工作簿的文件名列表
  csv_list (list): 输入CSV文件的路径列表
  output_path (string): Excel工作簿的输出路径
  """
  if check_lists_length(sheet_names, table_names, file_names, csv_list):
    raise Exception("CSV文件数、文件名、工作表名、表头名的数量必须一致")

  # 转换CSV路径为dbfs格式
  csv_list = ["dbfs:" + str.split(x,"dbfs")[1] for x in csv_list]
  # 转换输出路径为dbfs格式并创建目录
  path_api = "dbfs:" + str.split(output_path,"dbfs")[1]
  dbutils.fs.mkdirs(path_api)
  # 创建临时路径
  temp_path = "/tmp/" + str.split(dbutils.notebook.entry_point.getDbutils().notebook().getContext().userName().get(), sep = '@')[0]

  # 逐个处理CSV文件，避免内存堆积
  for i, (csv, sheet_name, table_name, file_name) in enumerate(zip(csv_list, sheet_names, table_names, file_names)):
    # 读取Spark DataFrame
    df_spark = spark.read.option("header", True) \
                .option("nullValue", "null") \
                .option("delimiter", ",") \
                .option("quote", "\"") \
                .option("escape", "\"") \
                .option("multiLine", True) \
                .option("inferSchema", True) \
                .option("enforceSchema", True) \
                .option("mode", "DROPMALFORMED") \
                .csv(csv)

    if df_spark.count() == 0:
      raise Exception(f'文件 {csv} 无有效行数据') 

    # 处理Decimal和Timestamp列
    dcml_cols = [f.name for f in df_spark.schema.fields if isinstance(f.dataType, DecimalType)]
    time_cols = [f.name for f in df_spark.schema.fields if isinstance(f.dataType, TimestampType)]
    for col_name in dcml_cols:
      df_spark = df_spark.withColumn(col_name, round(col(col_name).cast('double'),2))
    for col_name in time_cols:
      df_spark = df_spark.withColumn(col_name, to_date(col(col_name)))

    # 启用xlsxwriter常量内存模式，减少内存占用
    writer = pd.ExcelWriter(f"{temp_path}_{i}.xlsx", engine="xlsxwriter", options={"constant_memory": True})
    workbook = writer.book

    # 定义格式化样式（仅创建一次）
    top_border = workbook.add_format({'top': 2})
    bottom_border = workbook.add_format({'bottom': 2})
    left_border = workbook.add_format({'left': 2})
    right_border = workbook.add_format({'right': 2})
    top_left_border = workbook.add_format({'top': 2, 'left': 2})
    top_right_border = workbook.add_format({'top': 2, 'right': 2})
    bottom_left_border = workbook.add_format({'bottom': 2, 'left': 2})
    bottom_right_border = workbook.add_format({'bottom': 2, 'right': 2})
    table_title_format = workbook.add_format({'bold': True, 'font_color': 'white', 'bg_color': '#305496', 'align': 'left', 'valign': 'vcenter'})
    title_format = workbook.add_format({'bold': True, 'font_color': 'white', 'bg_color': 'black', 'align': 'left', 'valign': 'vcenter'})

    # 计算工作表拆分数量
    total_rows = df_spark.count()
    rows_per_sheet = 1000000
    num_sheets = (total_rows // rows_per_sheet) + 1 if total_rows % rows_per_sheet != 0 else total_rows // rows_per_sheet

    # 分批次写入Excel
    for j in range(num_sheets):
      start_row = j * rows_per_sheet
      # 读取当前批次数据
      df_part = df_spark.limit(rows_per_sheet).offset(start_row).toPandas()
      current_sheet_name = sheet_name + (f"_{j}" if j > 0 else "")
      
      # 写入数据到工作表
      df_part.to_excel(writer, sheet_name=current_sheet_name, index=False, header=False, startrow=4, startcol=1)
      worksheet = writer.sheets[current_sheet_name]

      # 写入表头
      for col_num, value in enumerate(df_part.columns.values):
        worksheet.write(3, col_num+1, value, table_title_format)

      num_part_rows = df_part.shape[0]
      num_cols = df_part.shape[1]

      # 设置边框格式
      worksheet.conditional_format(3,2,3,num_cols-1,{"type": "no_errors", "format": top_border})
      worksheet.conditional_format(4,1,num_part_rows+2,1,{"type": "no_errors", "format": left_border})
      worksheet.conditional_format(4,num_cols,num_part_rows+2,num_cols,{"type": "no_errors", "format": right_border})
      worksheet.conditional_format(num_part_rows+3,2,num_part_rows+3,num_cols-1,{"type": "no_errors", "format": bottom_border})
      worksheet.conditional_format(3,1,3,1,{"type": "no_errors", "format": top_left_border})
      worksheet.conditional_format(3,num_cols,3,num_cols,{"type": "no_errors", "format": top_right_border})
      worksheet.conditional_format(num_part_rows+3,1,num_part_rows+3,1,{"type": "no_errors", "format": bottom_left_border})
      worksheet.conditional_format(num_part_rows+3,num_cols,num_part_rows+3,num_cols,{"type": "no_errors", "format": bottom_right_border})

      # 设置工作表其他属性
      worksheet.hide_gridlines(2)
      worksheet.autofilter(3,1,num_part_rows+3,num_cols)
      worksheet.autofit()
      worksheet.set_column_pixels(0, 0, 20)
      worksheet.merge_range(1,1,1,num_cols, f"{table_name}_{j}" if j>0 else table_name, title_format)

      # 清理当前批次数据，释放内存
      del df_part
      gc.collect()

    # 关闭writer并移动文件到输出路径
    workbook.close()
    shutil.copy(f"{temp_path}_{i}.xlsx", output_path + file_name + ".xlsx")
    # 删除临时文件
    os.remove(f"{temp_path}_{i}.xlsx")
    print(f"文件 {file_name}.xlsx 创建成功")

    # 清理当前资源
    del df_spark, writer, workbook
    gc.collect()

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