Pyspark开发TF-IDF算法

113 阅读 0 评论 75 点赞

我是靠谱客的博主缓慢鲜花，这篇文章主要介绍Pyspark开发TF-IDF算法，现在分享给大家，希望可以做个参考。

直接上干货，一套操作猛如虎，一看结果很惊喜

复制代码

#! python3
# -*- coding: utf-8 -*-
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
spark = SparkSession.builder.appName("PySpark example").enableHiveSupport().getOrCreate()
spark.sparkContext.setLogLevel("WARN")
# Read data
df = spark.sql("SELECT img_label FROM sprs_log_basis.model_server_log WHERE datepart=20190425 LIMIT 10")
df.cache()
df.show()
# Get keys
df = df.select(F.map_keys("img_label").alias("keys"))
# Assign index
df = df.withColumn("doc_id", F.monotonically_increasing_id())
NUM_doc = df.count()
# One hot words
df = df.select('*', F.explode('keys').alias('token'))
df.show()
# Calculate TF
TF = df.groupBy("doc_id").agg(F.count("token").alias("doc_len")) 
.join(df.groupBy("doc_id", "token")
.agg(F.count("keys").alias("word_count")), ['doc_id']) 
.withColumn("tf", F.col("word_count") / F.col("doc_len")) 
.drop("doc_len", "word_count")
TF.cache()
# Calculate IDF
IDF = df.groupBy("token").agg(F.countDistinct("doc_id").alias("df"))
IDF = IDF.select('*', (F.log(NUM_doc / (IDF['df'] + 1))).alias('idf'))
IDF.cache()
# Calculate TF-IDF
TFIDF = TF.join(IDF, ['token']).withColumn('tf-idf', F.col('tf') * F.col('idf'))
TFIDF.show()
TFIDF.write.save("s3://***.tmp.ap-southeast-1/Default/hailin/here.csv", format='csv', header=True)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#! python3
# -*- coding: utf-8 -*-
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
spark = SparkSession.builder.appName("PySpark example").enableHiveSupport().getOrCreate()
spark.sparkContext.setLogLevel("WARN")
# Read data
df = spark.sql("SELECT img_label FROM sprs_log_basis.model_server_log WHERE datepart=20190425 LIMIT 10")
df.cache()
df.show()
# Get keys
df = df.select(F.map_keys("img_label").alias("keys"))
# Assign index
df = df.withColumn("doc_id", F.monotonically_increasing_id())
NUM_doc = df.count()
# One hot words
df = df.select('*', F.explode('keys').alias('token'))
df.show()
# Calculate TF
TF = df.groupBy("doc_id").agg(F.count("token").alias("doc_len")) 
.join(df.groupBy("doc_id", "token")
.agg(F.count("keys").alias("word_count")), ['doc_id']) 
.withColumn("tf", F.col("word_count") / F.col("doc_len")) 
.drop("doc_len", "word_count")
TF.cache()
# Calculate IDF
IDF = df.groupBy("token").agg(F.countDistinct("doc_id").alias("df"))
IDF = IDF.select('*', (F.log(NUM_doc / (IDF['df'] + 1))).alias('idf'))
IDF.cache()
# Calculate TF-IDF
TFIDF = TF.join(IDF, ['token']).withColumn('tf-idf', F.col('tf') * F.col('idf'))
TFIDF.show()
TFIDF.write.save("s3://***.tmp.ap-southeast-1/Default/hailin/here.csv", format='csv', header=True)