根据列的字符串值将数字列添加到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"}
...

我正在设计Naive Bayes分类器，这样的DataFrame是我的训练集：分类器将使用从字段 1 中提取的特征，而类（标签）由字段 type < / em>的

我的问题是我在拟合模型时遇到此错误：

  

pyspark.sql.utils.IllegalArgumentException：u'requirement failed：列类型必须是DoubleType类型，但实际上是StringType。'

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

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

我已经通过将StringIndexer用于DataFrame来解决。

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