spark中KMeans 文本聚类 对DataFrame中的文本计算tf-idf向量

在Spark中，可以使用`HashingTF`和`IDF`类来计算DataFrame中文本的tf-idf向量。然后，可以使用`KMeans`类对文本进行聚类。以下是一个简单的例子：

from pyspark.ml.clustering import KMeans
from pyspark.ml.feature import HashingTF, IDF, Tokenizer
from pyspark.ml.linalg import Vectors
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

# 创建一个DataFrame
df = spark.createDataFrame([(0, "hello world"), (1, "world is good"), (2, "goodbye world")], ["id", "sentence"])

# 分词
tokenizer = Tokenizer(inputCol="sentence", outputCol="words")
words_df = tokenizer.transform(df)

# 计算tf向量
hashing_tf = HashingTF(inputCol="words", outputCol="raw_features", numFeatures=20)
tf_df = hashing_tf.transform(words_df)

# 计算idf向量
idf = IDF(inputCol="raw_features", outputCol="tf_idf_features")
idf_model = idf.fit(tf_df)
tf_idf_df = idf_model.transform(tf_df)

# 提取tf-idf向量
tf_idf_df = tf_idf_df.select("id", "tf_idf_features")
to_vector = udf(lambda v: Vectors.dense(v.toArray()), returnType=VectorUDT())
tf_idf_df = tf_idf_df.withColumn("tf_idf_vector", to_vector("tf_idf_features"))

# 训练KMeans模型
kmeans = KMeans(k=2, seed=1)
model = kmeans.fit(tf_idf_df)

# 添加预测结果到DataFrame
predictions_df = model.transform(tf_idf_df)

# 展开tf-idf向量
for i in range(20):
  col_name = "tf_idf_" + str(i)
  tf_idf_df = tf_idf_df.withColumn(col_name, tf_idf_df.tf_idf_features.getItem(i))

# 删除原始的tf-idf向量列
tf_idf_df = tf_idf_df.drop("tf_idf_features")

# 合并预测结果和tf-idf向量
result_df = predictions_df.join(tf_idf_df, on="id")

这样，`result_df`就是包含原始文本、对应tf-idf向量和聚类结果的DataFrame了。