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

我正在设计一个朴素贝叶斯 classifier，这样的 DataFrame 是我的训练集：classifier 将使用从字段 1 中提取的特征class（标签）由字段 type.

给出

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

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)

Adding numeric column to pyspark DataFrame based on string value of column