使用 pandas_udf 和 Parquet 序列化时内存泄漏？
Posted
技术标签:
【中文标题】使用 pandas_udf 和 Parquet 序列化时内存泄漏？【英文标题】：Memory leaks when using pandas_udf and Parquet serialization? 【发布时间】：2019-10-13 04:59:42 【问题描述】：我目前正在使用 PySpark 开发我的第一个完整系统，但遇到了一些奇怪的与内存相关的问题。在其中一个阶段，我想类似于拆分-应用-组合策略来修改数据帧。也就是说，我想对给定列定义的每个组应用一个函数，最后将它们全部组合起来。问题是，我要应用的函数是“说出”熊猫习语的拟合模型的预测方法，即它是矢量化的，并将熊猫系列作为输入。
然后我设计了一个迭代策略，遍历组并手动应用 pandas_udf.Scalar 来解决问题。组合部分是使用对 DataFrame.unionByName() 的增量调用来完成的。我决定不使用 Pandas_udf 的 GroupedMap 类型，因为文档声明内存应由用户管理，并且当其中一个组可能太大而无法将其保存在内存中或由熊猫数据框。
主要问题是所有处理似乎都运行良好，但最后我想将最终的 DataFrame 序列化为 Parquet 文件。正是在这一点上，我收到了很多关于 DataFrameWriter 的类似 Java 的错误，或者内存不足的异常。
我已经在 Windows 和 Linux 机器上尝试过代码。我设法避免错误的唯一方法是增加机器中的 --driver-memory 值。每个平台的最小值都不同，并且取决于问题的大小，这让我怀疑是内存泄漏。
直到我开始使用 pandas_udf 才出现问题。我认为在使用 pandas_udf 时，pyarrow 序列化的整个过程中可能存在内存泄漏。
我创建了一个最小的可重现示例。如果我直接使用 Python 运行此脚本，则会产生错误。使用 spark-submit 并大量增加驱动内存，是可以让它工作的。
import pyspark
import pyspark.sql.functions as F
import pyspark.sql.types as spktyp
# Dummy pandas_udf -------------------------------------------------------------
@F.pandas_udf(spktyp.DoubleType())
def predict(x):
return x + 100.0
# Initialization ---------------------------------------------------------------
spark = pyspark.sql.SparkSession.builder.appName(
"mre").master("local[3]").getOrCreate()
sc = spark.sparkContext
# Generate a dataframe ---------------------------------------------------------
out_path = "out.parquet"
z = 105
m = 750000
schema = spktyp.StructType(
[spktyp.StructField("ID", spktyp.DoubleType(), True)]
)
df = spark.createDataFrame(
[(float(i),) for i in range(m)],
schema
)
for j in range(z):
df = df.withColumn(
f"Nj",
F.col("ID") + float(j)
)
df = df.withColumn(
"X",
F.array(
F.lit("A"),
F.lit("B"),
F.lit("C"),
F.lit("D"),
F.lit("E")
).getItem(
(F.rand()*3).cast("int")
)
)
# Set the column names for grouping, input and output --------------------------
group_col = "X"
in_col = "N0"
out_col = "EP"
# Extract different group ids in grouping variable -----------------------------
rows = df.select(group_col).distinct().collect()
groups = [row[group_col] for row in rows]
print(f"Groups: groups")
# Split and treat the first id -------------------------------------------------
first, *others = groups
cur_df = df.filter(F.col(group_col) == first)
result = cur_df.withColumn(
out_col,
predict(in_col)
)
# Traverse the remaining group ids ---------------------------------------------
for i, other in enumerate(others):
cur_df = df.filter(F.col(group_col) == other)
new_df = cur_df.withColumn(
out_col,
predict(in_col)
)
# Incremental union --------------------------------------------------------
result = result.unionByName(new_df)
# Save to disk -----------------------------------------------------------------
result.write.mode("overwrite").parquet(out_path)
令人震惊（至少对我而言），如果我在序列化语句之前调用 repartition()，问题似乎就消失了。
result = result.repartition(result.rdd.getNumPartitions())
result.write.mode("overwrite").parquet(out_path)
将这一行放在适当的位置后，我可以降低很多驱动程序内存配置，并且脚本运行良好。我几乎无法理解所有这些因素之间的关系，尽管我怀疑代码的惰性评估和 pyarrow 序列化可能是相关的。
这是我目前用于开发的环境：
arrow-cpp 0.13.0 py36hee3af98_1 conda-forge
asn1crypto 0.24.0 py36_1003 conda-forge
astroid 2.2.5 py36_0
atomicwrites 1.3.0 py_0 conda-forge
attrs 19.1.0 py_0 conda-forge
blas 1.0 mkl
boost-cpp 1.68.0 h6a4c333_1000 conda-forge
brotli 1.0.7 he025d50_1000 conda-forge
ca-certificates 2019.3.9 hecc5488_0 conda-forge
certifi 2019.3.9 py36_0 conda-forge
cffi 1.12.3 py36hb32ad35_0 conda-forge
chardet 3.0.4 py36_1003 conda-forge
colorama 0.4.1 py36_0
cryptography 2.6.1 py36hb32ad35_0 conda-forge
dill 0.2.9 py36_0
docopt 0.6.2 py36_0
entrypoints 0.3 py36_0
falcon 1.4.1.post1 py36hfa6e2cd_1000 conda-forge
fastavro 0.21.21 py36hfa6e2cd_0 conda-forge
flake8 3.7.7 py36_0
future 0.17.1 py36_1000 conda-forge
gflags 2.2.2 ha925a31_0
glog 0.3.5 h6538335_1
hug 2.5.2 py36hfa6e2cd_0 conda-forge
icc_rt 2019.0.0 h0cc432a_1
idna 2.8 py36_1000 conda-forge
intel-openmp 2019.3 203
isort 4.3.17 py36_0
lazy-object-proxy 1.3.1 py36hfa6e2cd_2
libboost 1.67.0 hd9e427e_4
libprotobuf 3.7.1 h1a1b453_0 conda-forge
lz4-c 1.8.1.2 h2fa13f4_0
mccabe 0.6.1 py36_1
mkl 2018.0.3 1
mkl_fft 1.0.6 py36hdbbee80_0
mkl_random 1.0.1 py36h77b88f5_1
more-itertools 4.3.0 py36_1000 conda-forge
ninabrlong 0.1.0 dev_0 <develop>
nose 1.3.7 py36_1002 conda-forge
nose-exclude 0.5.0 py_0 conda-forge
numpy 1.15.0 py36h9fa60d3_0
numpy-base 1.15.0 py36h4a99626_0
openssl 1.1.1b hfa6e2cd_2 conda-forge
pandas 0.23.3 py36h830ac7b_0
parquet-cpp 1.5.1 2 conda-forge
pip 19.0.3 py36_0
pluggy 0.11.0 py_0 conda-forge
progressbar2 3.38.0 py_1 conda-forge
py 1.8.0 py_0 conda-forge
py4j 0.10.7 py36_0
pyarrow 0.13.0 py36h8c67754_0 conda-forge
pycodestyle 2.5.0 py36_0
pycparser 2.19 py36_1 conda-forge
pyflakes 2.1.1 py36_0
pygam 0.8.0 py_0 conda-forge
pylint 2.3.1 py36_0
pyopenssl 19.0.0 py36_0 conda-forge
pyreadline 2.1 py36_1
pysocks 1.6.8 py36_1002 conda-forge
pyspark 2.4.1 py_0
pytest 4.5.0 py36_0 conda-forge
pytest-runner 4.4 py_0 conda-forge
python 3.6.6 hea74fb7_0
python-dateutil 2.8.0 py36_0
python-hdfs 2.3.1 py_0 conda-forge
python-mimeparse 1.6.0 py_1 conda-forge
python-utils 2.3.0 py_1 conda-forge
pytz 2019.1 py_0
re2 2019.04.01 vc14h6538335_0 [vc14] conda-forge
requests 2.21.0 py36_1000 conda-forge
requests-kerberos 0.12.0 py36_0
scikit-learn 0.20.1 py36hb854c30_0
scipy 1.1.0 py36hc28095f_0
setuptools 41.0.0 py36_0
six 1.12.0 py36_0
snappy 1.1.7 h777316e_3
sqlite 3.28.0 he774522_0
thrift-cpp 0.12.0 h59828bf_1002 conda-forge
typed-ast 1.3.1 py36he774522_0
urllib3 1.24.2 py36_0 conda-forge
vc 14.1 h0510ff6_4
vs2015_runtime 14.15.26706 h3a45250_0
wcwidth 0.1.7 py_1 conda-forge
wheel 0.33.1 py36_0
win_inet_pton 1.1.0 py36_0 conda-forge
wincertstore 0.2 py36h7fe50ca_0
winkerberos 0.7.0 py36_1
wrapt 1.11.1 py36he774522_0
xz 5.2.4 h2fa13f4_4
zlib 1.2.11 h62dcd97_3
zstd 1.3.3 hfe6a214_0
任何提示或帮助将不胜感激。
【问题讨论】：
你能在初始化数据框后立即重新分区看看吗？您还可以提供数据大小和失败/通过的内存配置。这肯定会有所帮助 嗨。谢谢你的评论。刚从 Parquet 文件中读取数据后，我一开始就尝试重新分区，但它不起作用。我的猜测是重新分区似乎修复了在 udf 阶段损坏的东西。关于数据大小，似乎与您运行示例的机器有关。在我的 Windows 笔记本电脑 8GB RAM 下破坏它（不使用重新分区时）比在我用作测试环境的 Ubuntu VM 中更容易。如果可能，我会尝试提供一些数字并适当地编辑问题。谢谢。 @Fernandez 我遇到了类似的问题，想知道您是否找到任何解决方案。我尝试在将 df 写入镶木地板之前对其进行重新分区，但这并没有像您的情况那样有帮助。 【参考方案1】：我想对你的帖子发表评论，但我的声誉太低了。
根据我的经验，udf 会大大降低你的性能，特别是如果你用 python（或 pandas？）编写它们。有一篇文章，为什么你不应该使用python udfs而使用scala udfs：https://medium.com/wbaa/using-scala-udfs-in-pyspark-b70033dd69b9
在我的情况下，可以使用内置函数，即使它非常复杂，运行时间也比以前减少了大约 5%。
对于您的 OOM 错误以及为什么重新分区对您有效，我无法解释。我能给你的唯一建议是尽可能避免使用 UDF，尽管在你的情况下似乎并不那么容易。
【讨论】：
这篇文章很有意思。谢谢你。问题是，我在这里尝试使用的是 pandas_udf，而不是普通的 UDF。我们之所以需要它，是因为我们想利用类 sklearn 模型的矢量化 predict() 方法。使用 UDF 意味着我们必须每次观察调用一次模型，我想避免这种情况。【参考方案2】：这个帖子有点老了，但我遇到了完全相同的问题并花了好几个小时来解决它。所以我只是想解释一下我是如何解决它的，希望它能为以后遇到同样问题的其他人节省一些时间。
这里的问题与pandas_udf或parquet无关，而是使用withColumn生成列。当向数据框添加多列时，使用select 方法更有效。 This article 解释了原因。
例如，而不是
for j in range(z):
df = df.withColumn(
f"Nj",
F.col("ID") + float(j)
)
你应该写
df = df.select(
*df.columns,
*[(F.col("ID") + float(j)).alias(f"Nj") for j in range(z)]
)
重写后的脚本如下所示（请注意，我仍然必须将驱动程序内存增加到 2GB，但至少是相当合理的内存量）
import pyspark
import pyspark.sql.functions as F
import pyspark.sql.types as spktyp
# Dummy pandas_udf -------------------------------------------------------------
@F.pandas_udf(spktyp.DoubleType())
def predict(x):
return x + 100.0
# Initialization ---------------------------------------------------------------
spark = (pyspark.sql.SparkSession.builder
.appName("mre")
.config("spark.driver.memory", "2g")
.master("local[3]").getOrCreate())
sc = spark.sparkContext
# Generate a dataframe ---------------------------------------------------------
out_path = "out.parquet"
z = 105
m = 750000
schema = spktyp.StructType(
[spktyp.StructField("ID", spktyp.DoubleType(), True)]
)
df = spark.createDataFrame(
[(float(i),) for i in range(m)],
schema
)
df = df.select(
*df.columns,
*[(F.col("ID") + float(j)).alias(f"Nj") for j in range(z)]
)
df = df.withColumn(
"X",
F.array(
F.lit("A"),
F.lit("B"),
F.lit("C"),
F.lit("D"),
F.lit("E")
).getItem(
(F.rand()*3).cast("int")
)
)
# Set the column names for grouping, input and output --------------------------
group_col = "X"
in_col = "N0"
out_col = "EP"
# Extract different group ids in grouping variable -----------------------------
rows = df.select(group_col).distinct().collect()
groups = [row[group_col] for row in rows]
print(f"Groups: groups")
# Split and treat the first id -------------------------------------------------
first, *others = groups
cur_df = df.filter(F.col(group_col) == first)
result = cur_df.withColumn(
out_col,
predict(in_col)
)
# Traverse the remaining group ids ---------------------------------------------
for i, other in enumerate(others):
cur_df = df.filter(F.col(group_col) == other)
new_df = cur_df.select(
*cur_df.columns,
predict(in_col).alias(out_col)
)
# Incremental union --------------------------------------------------------
result = result.unionByName(new_df)
# Save to disk -----------------------------------------------------------------
result.write.mode("overwrite").parquet(out_path)
【讨论】：
我面临着类似的问题。在我的情况下，虽然我只有一列添加了“withColumn”，所以使用 select 对我的情况没有帮助。还有其他可能对您有帮助的线索吗？谢谢！以上是关于使用 pandas_udf 和 Parquet 序列化时内存泄漏？的主要内容，如果未能解决你的问题，请参考以下文章
为啥我的应用程序不能以 pandas_udf 和 PySpark+Flask 开头？