Pyspark系列笔记--如何在一个pysprk Dataframe上训练word2vec模型

前言

本次试验环境：

spark-1.5.0
python-2.7

---

Step 1. Create a dataframe use a json file

值得注意的事
json格式要求每一行为一个json串，即json串不能跨行。支持json格式处理的有jq插件。可参考这一篇关于jq插件的博客。

print('读取json文件...')
from pyspark.sql import SQLContext
sqlContext = SQLContext(sc)
path = r'hdfs://master:9000/test.json'
jsonDF = sqlContext.read.json(path).cache()
jsonDf.show()

---

Step 2. Transform the column you want to use

第一步之后，我们就已经创建了一个dataframe。
show函数以后，大致会显示出以下的数据结构：

DataFrame结构图

假设我们要对name这一列的名字做word2vec的模型，我们首先得要对这一列进行分词。

在单机的python中我们很容易做到，那么在pyspark中就没那么简单了。
我们首先定义一个分词函数：

import re
from pyspark.sql import Row
def normalize_name(name):
    name = re.sub(r"[.,\/#!$%\^\*;:{}=\_`~()@]", ' ', name)
    name = re.sub(r'\s+', ' ', name).strip()
    # name = [tok for tok in re.split(" ",name) if tok]
    name = name.split(" ")
    return name

因为这里的数据比较简单，我们用空格切分。
为了使得该函数应用在该列的每个元素上，我们可以用map，

jsonDF.name= jsonDF.select("name").map(boolMap)

但是map以后的数据并不是dataframe的列格式，导致修改失败。（并没有试过是否可以直接将map以后的结果强转为Column，下次可以试一下）。
最好的方式是将该函数转成udf的函数形式：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType,StringType
normal = udf(normalize_name,ArrayType(StringType()))
jsonDF = jsonDF.withColumn("name",normal(jsonDF['name']))

需要特意强调的是这里的udf中的第二个参数，因为word2vec要求传入的列是ArrayType的类型，我当时候调了好久。ArrayType有点像元组格式，其具体形式为([a,b,c],True)，abc是切分以后的词。当时候在函数返回里面直接返回tuple是不对的，返回的是惰性计算的结果，如果在withColumn里面直接转的话，又要求是一个Column类型，所以必须显式指定返回类型，就可以自动将返回类型包装啦。

然后我们可以愉快的训练word2vec了。

---

Step 3. Training Word2Vec

from pyspark.ml.feature import Word2Vec
model = Word2Vec(vectorSize=5, seed=42, inputCol="name", outputCol="model").fit(jsonDF)
model.getVectors().show()

注意这里的inputCol的列名是你要做词向量的列。

这样我们就能得到词向量啦：

最终词向量结果