使用pyspark.sql.GroupedData.applyInPandas出现非预期原地修改对象问题

在使用pyspark.sql.GroupedData.applyInPandas和pandas.core.groupby.GroupBy.apply执行同一函数时，出现了不一致的结果：Spark会意外地原地修改外部传入的对象，而Pandas则不会。

最小复现示例 #

import numpy as np
import pandas as pd
from pyspark.sql import SparkSession

spark = SparkSession \
    .builder \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .config("spark.sql.execution.arrow.pyspark.fallback.enabled", "true") \
    .getOrCreate()

def func(df_1, df_2):
    # do something with df_1
    # ...
    
    # this shouldn't modify df_2
    print("df_2 before", df_2.values[0], end="\t")
    mask = np.array([True])
    out = df_2["col_1"].astype(str)
    out[mask] = out[mask].astype(float) / 10
    print("df_2 after", df_2.values[0])

    return pd.DataFrame({"out": [out.astype(str).iloc[0]]})


# sample data
second_df = pd.DataFrame({"col_1": ["90"]})

pdf = pd.DataFrame({"x":['0', '1', '2', '3']})
sdf = spark.createDataFrame(pdf)

Pandas 执行结果 #

使用Pandas调用func()时，外部对象second_df并未被修改：

_ = pdf.groupby("x").apply(lambda x: func(x, second_df))

#    df_2 before ['90'] df_2 after ['90']
#    df_2 before ['90'] df_2 after ['90']
#    df_2 before ['90'] df_2 after ['90']
#    df_2 before ['90'] df_2 after ['90']

PySpark 执行结果 #

但使用PySpark调用时，second_df在函数作用域内被持续修改：

_ = sdf.groupby("x").applyInPandas(lambda x : func(x, second_df), schema="out string").toPandas()

#    df_2 before ['90'] df_2 after [9.0]
#    df_2 before [9.0]  df_2 after [0.9]
#    df_2 before [0.9]  df_2 after [0.09]
#    df_2 before [0.09] df_2 after [0.009]

临时解决方法 #

目前找到的有效解决方法是将外部对象的定义移至函数内部，确保每次调用函数时都使用全新的对象：

def func2(x):
    # sample data
    second_df = pd.DataFrame({"col_1": ["90"]})
    return func(x, second_df)

_ = pdf.groupby("x").apply(func2)

_ = sdf.groupby("x").applyInPandas(func2, schema="out string").toPandas()

# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']
# df_2 before ['90']      df_2 after ['90']

环境信息 #

python = 3.8.9
pandas = 1.5.3
numpy = 1.24.2
pyspark = 3.3.1
pyarrow = 11.0.0

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