PySpark导入MySQL数据后行丢失问题：原因及解决方法咨询

问题背景 #

任务流程：

1. 从MySQL导入数据到PySpark
2. 执行数据转换操作
3. 将转换后的数据写回MySQL数据库

核心代码框架 #

# load the SparkSession 
configs = get_configs_for_spark()
spark = get_spark_session(configs)

# open SSH tunnel with VM 
with SSHTunnelForwarder(**get_ssh_tunnel_args(configs)) as tunnel: 
    # grab unprocessed data
    df = get_raw_data(configs, tunnel, spark)

    # transform the data 
    df = transform_my_data_1(configs, df)

    # Get the number of rows in the dataframe
    num_rows = df.count()

    # Determine the number of subsets (each with 100 rows)
    subset_size = 100
    num_partitions = int(np.ceil(num_rows / subset_size))

    # Add a partition_id column to your DataFrame
    df = df.repartition(num_partitions)  # Specify the desired number of partitions

    # Add a partition_id column to your DataFrame
    df = df.withColumn("partition_id", spark_partition_id())

    # Group by partition_id and count the occurrences
    partition_counts = df.groupBy("partition_id").count()

    # create the Azure client to store data to blob storage
    container_client = create_azure_client(configs)  


    # Iterate over the subsets
    for i in range(num_partitions):
        subset_df = df.filter(df.partition_id == i)

        # Run transform_my_data_2 on the subset dataframe
        subset_df = transform_my_data_2(configs, container_client, subset_df)

        # Write the subset dataframe to MySQL DB
        write_data_to_mysql_db(configs, tunnel, subset_df.drop("partition_id"))

        # print the size of the data frame and expected size 
        # Print progress when partitions have been executed
        expected_df_count = partition_counts.filter(partition_counts.partition_id == i).select("count").first()[0]
        print(f"{i+1} / {num_partitions} partitions processed. || subset df Size: Expected({expected_df_count}) Actual({subset_df.count()}) ")

    
    # end spark session 
    spark.stop()

MySQL读写函数 #

写入函数

def write_data_to_mysql_db(configs, tunnel, df):
    database = configs["MySQL"]["database"]
    username = configs["MySQL"]["username"]
    password = configs["MySQL"]["password"]
    driver = configs["MySQL"]["driver"]
    table = configs["MySQL"]["table"]

    # Define the JDBC URL for the MySQL database
    url = f"jdbc:mysql://localhost:{tunnel.local_bind_port}/{database}"
    # Define the database properties for authentication
    properties = {
        "user": username,
        "password": password,
        "driver": driver
    }
    df.write.jdbc(url=url, table=table, mode="append", properties=properties)

读取函数

def get_raw_data(configs, tunnel, spark_session):
    """
    gets raw data
    """
    database = configs["MySQL"]["database"]
    username = configs["MySQL"]["username"]
    password = configs["MySQL"]["password"]

    # Define the MySQL JDBC URL
    url = f'jdbc:mysql://localhost:{tunnel.local_bind_port}/{database}'
   
    # Create a dataframe by reading the data from the MySQL database
    df = spark_session.read.jdbc(url=url, table=transription_table, 
                         properties={"user": username, 
                                     "password": password,
                                     "characterSetResults": "utf8mb4"})


    # Execute custom SQL query to get data
    my_query = f"""QUERY REDACTED"""

    df = spark_session.read.jdbc(url=url, table=f"({my_query}) as temp", properties={"user": username, "password": password})


    return df

问题现象 #

• 读取数据并完成第一次转换后，数据行数符合预期
• 完整脚本运行后出现部分行丢失，分区统计前后不一致：
  初始分区统计：

  +------------+-----+
  |partition_id|count|
  +------------+-----+
  |           0|   90|
  |           1|   90|
  |           2|   90|
  |           3|   90|
  |           4|   90|
  |           5|   89|
  |           6|   89|
  |           7|   90|
  |           8|   90|
  +------------+-----+
  

  多次写入后分区统计：

  +------------+-----+
  |partition_id|count|
  +------------+-----+
  |           0|   87|
  |           1|   87|
  |           2|   87|
  |           3|   87|
  |           4|   87|
  |           5|   86|
  |           6|   86|
  |           7|   87|
  |           8|   87|
  +------------+-----+
  

• 最终MySQL数据库中数据缺失

已尝试的调试方法 #

调试1：模拟数据测试

使用脚本生成的模拟数据运行时，分区行数始终稳定，无行丢失：

+------------+-----+
|partition_id|count|
+------------+-----+
|           0|  100|
|           1|  100|
|           2|  100|
|           3|  100|
|           4|  100|
|           5|  100|
|           6|  100|
|           7|  100|
|           8|  100|
|           9|  100|
+------------+-----+

调试2：使用df.persist()

该方法解决了行丢失问题，但计算速度极慢，仅能处理极小数据集：

# Iterate over the subsets
for i in range(num_partitions):
    if not df.is_cached():
        df.persist() 

    subset_df = df.filter(df.partition_id == i)

    # Run transform_my_data_2 on the subset dataframe
    subset_df = transform_my_data_2(configs, container_client, subset_df)

    # Write the subset dataframe to MySQL DB
    write_data_to_mysql_db(configs, tunnel, subset_df.drop("partition_id"))

    subset_df.unpersist()

    # print the size of the data frame and expected size 
    # Print progress when partitions have been executed
    expected_df_count = partition_counts.filter(partition_counts.partition_id == i).select("count").first()[0]
    print(f"{i+1} / {num_partitions} partitions processed. || subset df Size: Expected({expected_df_count}) Actual({subset_df.count()}) ")
df.unpersist()

疑问 #

1. 为何模拟数据正常，MySQL导入数据异常？
2. df.persist()解决问题但速度慢的原因？
3. 如何在真实数据上避免行丢失？

解答 #

1. 模拟数据正常、MySQL导入数据异常的原因 #

• 数据源特性差异：模拟数据是内存生成的确定性数据，而MySQL导入的数据可能包含非确定性元素（如NULL值、特殊字符，或查询中使用了RAND()、NOW()这类非确定性函数）。Spark重新计算DataFrame时，非确定性操作会导致每次生成的分区数据不一致，进而出现行丢失。
• 分区逻辑的非确定性：repartition(num_partitions)默认用哈希分区，如果数据中存在大量NULL值，所有NULL会被分到同一分区；后续多次计算时，spark_partition_id()返回的结果可能因分区数据重新分布而变化，导致过滤后的行数不一致。
• JDBC读取的隐式分区：通过JDBC从MySQL读取数据时，Spark会将数据拆分成多个任务读取，任务执行顺序或数据分片存在不确定性，重新计算DataFrame时，数据的分区分布会发生变化。

2. df.persist()解决问题但速度慢的原因 #

• persist()会把DataFrame的计算结果缓存到内存或磁盘，避免每次循环都重新执行从MySQL读取、转换、分区的全量计算，因此解决了行丢失问题（缓存的是固定计算结果，不会重新生成分区数据）。
• 速度慢的核心原因：
  • 默认存储级别是MEMORY_ONLY，数据量较大时内存无法容纳，会溢出到磁盘，磁盘IO速度远低于内存，拖慢整体计算。
  • 循环内重复执行缓存检查和操作，带来额外开销（实际上只需要在循环前缓存一次即可）。
  • MySQL JDBC写入默认是单线程操作，加上循环次数多，进一步放大了写入的开销。

3. 真实数据避免行丢失的优化方案 #

方案1：使用确定性分区方式替代spark_partition_id()

不要依赖spark_partition_id()这种非确定性标识，改用固定逻辑生成分区ID：

from pyspark.sql.functions import monotonically_increasing_id, floor

# 生成唯一递增ID，按子集大小计算固定分区ID
df = df.withColumn("row_id", monotonically_increasing_id())
df = df.withColumn("partition_id", floor(df["row_id"] / subset_size))

# 统计各分区行数
partition_counts = df.groupBy("partition_id").count()

这种方式生成的partition_id是固定的，不会因重新计算而变化，彻底避免分区数据不一致。

方案2：提前物化分区数据

生成partition_id后，将数据写入临时视图，后续循环直接读取物化后的结果：

# 将分区后的数据写入临时视图
df.createOrReplaceTempView("partitioned_data")

# 循环读取每个分区
for i in range(num_partitions):
    subset_df = spark.sql(f"SELECT * FROM partitioned_data WHERE partition_id = {i}")
    # 执行转换和写入操作...

临时视图的数据是物化后的固定结果，不会重复执行上游的读取和转换逻辑。

方案3：优化persist()存储级别

如果必须用缓存，选择更高效的存储级别，比如MEMORY_AND_DISK_SER（序列化后存储在内存+磁盘），减少内存占用并提升IO效率：

import pyspark

# 循环前仅缓存一次
df.persist(pyspark.StorageLevel.MEMORY_AND_DISK_SER)

for i in range(num_partitions):
    subset_df = df.filter(df.partition_id == i)
    # 执行转换和写入操作...

df.unpersist()

方案4：减少循环次数，批量写入

增大分区大小（比如从100行改为10000行），减少循环次数，降低JDBC写入的开销；或者直接使用Spark原生批量写入能力，无需手动分区：

# 去掉手动分区和循环，直接全量转换后写入
df = transform_my_data_2(configs, container_client, df)
write_data_to_mysql_db(configs, tunnel, df)

如果必须分批写入，可使用foreachBatch自动处理分区：

def write_batch(batch_df, batch_id):
    write_data_to_mysql_db(configs, tunnel, batch_df)

df.writeStream.foreachBatch(write_batch).start().awaitTermination()

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