在 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)

现在，如果我尝试手动索引这些值并将它们存储在数据框中，一切都会像魅力一样运作，所有这些都在几个亚马逊 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 在 "user" 列上调用 countByValue 以获取所有不同的值。

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

然而，随着不同值的数量增加，问题只会再次发生。不幸的是，您无法使用 Spark 的转换器做很多事情，原因如下。实际上，在对数据进行拟合之后，transformer 应该被序列化以供进一步使用。因此它们并没有被设计成这么大（一个 15M 字符串的地图至少有 100MB 的重量）。我认为您需要重新考虑对那么多类别使用 StringIndexer。使用散列技巧可能更适合这里。

最后，让我评论一下您的解决方法。使用 window，您实际上将所有 15M 类别放在一个分区上，因此放在一个执行器上。如果该数字增加，它将不会扩展。此外，使用非分区 window 通常不是一个好主意，因为它会阻止并行计算（除了将所有内容放在同一分区上可能会导致 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"])

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

In Spark ML, why is fitting a StringIndexer on a column with million of disctinct values yielding an OOM error?

apache-spark

pyspark

apache-spark-ml