在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 不同值，您实际上是在驱动程序上创建一个巨大的映射，并且它耗尽了内存......一个简单的解决方法是增加驱动程序的内存。如果您使用spark-submit，您可以使用--driver-memory 16g. 您还可以使用spark.driver.memory配置文件中的属性。

然而，随着不同值数量的增加，这个问题将会再次出现。不幸的是，您无法使用 Spark 的 Transformer 做太多事情，原因如下。实际上，在适合数据之后，变压器应该被序列化以供进一步使用。因此它们并没有设计得这么大（具有 15M 字符串的地图至少有 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"])