python - 根据列的字符串值将数字列添加到 pyspark DataFrame

Question

我有一个从 JSON 文件构建的 DataFrame:

{ "1": "a b c d e f", "2": 1, "type": "type1"}
{ "1": "a b c b c", "2": 2, "type": "type1"}
{"1": "d d a b c", "2": 3, "type": "type2"}
...

我正在设计一个朴素贝叶斯分类器，这样的 DataFrame 就是我的训练集:分类器将使用从字段 1 中提取的特征，类(标签)由字段 type< 给出/em>.

我的问题是在拟合模型时出现此错误:

pyspark.sql.utils.IllegalArgumentException: u'requirement failed: Column type must be of type DoubleType but was actually StringType.'

表示标签字段必须是数字。为了解决这个问题，我试图通过字典将字符串值映射为数值

grouped = df.groupBy(df.type).agg({'*': 'count'}).persist()
types = {row.type: grouped.collect().index(row) for row in grouped.collect()}

然后我们的想法是向 DataFrame 添加一个新列，其数值对应于它的字符串值:

df = df.withColumn('type_numeric', types[df.type])

这当然失败了，所以我想知道是否有人对如何实现这一目标有更好的想法或建议。

Answer 1

我已经通过使用 DataFrame 的 StringIndexer 解决了问题。

string_indexer = StringIndexer(inputCol='type', outputCol='type_numeric')
rescaled_data_numeric = string_indexer.fit(df).transform(df)