apache-spark - 在 Spark ML 中，为什么在具有数百万个不同值的列上安装 StringIndexer 会产生 OOM 错误？

Question

我正在尝试在具有大约 15.000.000 个唯一字符串值的列上使用 Spark 的 StringIndexer 特征转换器。无论我投入多少资源，Spark 总是会因某种内存不足异常而死在我身上。

from pyspark.ml.feature import StringIndexer

data = spark.read.parquet("s3://example/data-raw").select("user", "count")

user_indexer = StringIndexer(inputCol="user", outputCol="user_idx")

indexer_model = user_indexer.fit(data) # This never finishes

indexer_model \
    .transform(data) \
    .write.parquet("s3://example/data-indexed")

驱动程序上会生成一个错误文件，其开头如下所示:

#
# There is insufficient memory for the Java Runtime Environment to continue.
# Native memory allocation (mmap) failed to map 268435456 bytes for committing reserved memory.
# Possible reasons:
#   The system is out of physical RAM or swap space
#   In 32 bit mode, the process size limit was hit
# Possible solutions:
#   Reduce memory load on the system
#   Increase physical memory or swap space
#   Check if swap backing store is full
#   Use 64 bit Java on a 64 bit OS
#   Decrease Java heap size (-Xmx/-Xms)
#   Decrease number of Java threads
#   Decrease Java thread stack sizes (-Xss)
#   Set larger code cache with -XX:ReservedCodeCacheSize=
# This output file may be truncated or incomplete.
#
#  Out of Memory Error (os_linux.cpp:2657)

现在，如果我尝试手动索引值并将它们存储在数据帧中，一切都会像魅力一样工作，所有这些都在几个 Amazon c3.2xlarge 工作人员上进行。

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

data = spark.read.parquet("s3://example/data-raw").select("user", "count")

uid_map = data \
    .select("user") \
    .distinct() \
    .select("user", row_number().over(Window.orderBy("user")).alias("user_idx"))

data.join(uid_map, "user", "inner").write.parquet("s3://example/data-indexed")

我真的很想使用 Spark 提供的正式转换器，但目前这似乎不可能。关于如何完成这项工作的任何想法？

Answer 1

你得到 OOM 错误的原因是在幕后，Spark 的 StringIndexer电话countByValue在“用户”列上获取所有不同的值。

使用 15M 不同的值，您实际上是在驱动程序上创建一个巨大的 Map 并且它耗尽了内存......一个简单的解决方法是增加驱动程序的内存。如果你使用 spark-submit 你可以使用 --driver-memory 16g .您也可以使用 spark.driver.memory配置文件中的属性。

然而，随着不同值数量的增加，问题将再次发生。不幸的是，Spark 的转换器无能为力，这就是原因。实际上，在适应数据之后，转换器旨在被序列化以供进一步使用。因此，它们没有被设计成这么大(一个包含 15M 字符串的 map 至少会重 100MB)。我认为您需要重新考虑对这么多类别使用 StringIndexer。在这里使用哈希技巧可能更适合。

最后，让我评论一下您的解决方法。使用您的窗口，您实际上将所有 15M 类别放在一个分区上，因此放在一个执行程序上。如果这个数字增加，它就不会扩展。此外，使用非分区窗口通常是一个坏主意，因为它会阻止并行计算(除了将所有内容都放在同一个分区上会导致 OOM 错误)。我会计算你的 uid_map像这样:

# if you don't need consecutive indices
uid_map = data\
    .select("user")\
    .distinct()\
    .withColumn("user_idx", monotonically_increasing_id())

# if you do, you need to use RDDs
uid_rdd = data\
    .select("user")\
    .distinct()\
    .rdd.map(lambda x : x["user"])\
    .zipWithIndex()
uid_map = spark.createDataFrame(uid_rdd, ["user", "user_idx"])