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

Question

我试图在一个具有大约15.000.000个唯一字符串值的列上使用Spark的StringIndexer特征变换器。无论我投入多少资源，Spark都会因为某种Out Of Memory异常而死在我身上。

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)

现在，如果我尝试手动索引值并将它们存储在数据框中，那么一切都像魅力一样，都在几个亚马逊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字符串的地图至少会重100MB）。我认为您需要重新考虑使用StringIndexer来处理这么多类别。使用Hashing技巧可能更适合这里。

最后，让我评论一下你的解决方法。通过窗口，您实际上将所有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"])