在 pyspark 中用整数对列进行编码

Question

我必须在 pyspark(spark 2.0) 的大 DataFrame 中对列进行编码。所有的值几乎都是唯一的（大约10亿个值）。最好的选择可能是 StringIndexer，但由于某种原因它总是失败并终止我的 Spark 会话。我可以以某种方式编写这样的函数吗：

id_dict() = dict()
def indexer(x):
    id_dict.setdefault(x, len(id_dict))
    return id_dict[x]

并使用 id_dict 将其映射到 DataFrame 保存 items()？这个字典会在每个执行器上同步吗？我需要所有这些来预处理 Spark.mllib ALS 模型的元组 ('x', 3, 5)。谢谢。

Answer 1

StringIndexer将所有标签保留在内存中，因此如果值几乎是唯一的，它就无法扩展。

您可以采用唯一值，排序并添加 id，这很昂贵，但在这种情况下更强大：

from pyspark.sql.functions import monotonically_increasing_id

df = spark.createDataFrame(["a", "b", "c", "a", "d"], "string").toDF("value")

indexer = (df.select("value").distinct()
  .orderBy("value")
  .withColumn("label", monotonically_increasing_id()))

df.join(indexer, ["value"]).show()
# +-----+-----------+
# |value|      label|
# +-----+-----------+
# |    d|25769803776|
# |    c|17179869184|
# |    b| 8589934592|
# |    a|          0|
# |    a|          0|
# +-----+-----------+

请注意，标签不是连续的，并且每次运行可能会有所不同，或者如果发生spark.sql.shuffle.partitions变化也可能会发生变化。如果不可接受，您将不得不使用RDDs：

from operator import itemgetter

indexer = (df.select("value").distinct()
    .rdd.map(itemgetter(0)).zipWithIndex()
    .toDF(["value", "label"]))

df.join(indexer, ["value"]).show()
# +-----+-----+
# |value|label|
# +-----+-----+
# |    d|    0|
# |    c|    1|
# |    b|    2|
# |    a|    3|
# |    a|    3|
# +-----+-----+