PySpark TF-IDF計算（2）

使用PySpark進行TF-IDF計算

這篇博文將記錄使用PySpark進行TF-IDF統計的過程，將提供多種計算方法。

1. 準備數據

爲了簡單，同時爲了驗證自己的程序有木有錯誤，我使用如下的測試數據：

1 我來到北京清華大學
2 他來到了網易杭研大廈
3 我來到北京清華大學
4 他來到了網易杭研大廈
5 我來到北京清華大學，我來到北京清華大學

一共五行，每行代表一篇文章，每行中得文章id和正文使用空格分開，例如第一行：1代表文章id,"我來到北京清華大學"代表一篇文本。
將文本寫入到文件test中。

2. 加載數據並且分詞

分詞采用jieba分詞，代碼如下：

def seg(data):
    """
    分詞後返回分詞的dataframe
    :param spark:
    :param data:
    :return:
    """
    return [w for w in jieba.cut(data.strip(), cut_all=False) if len(w) > 1 and re.match(remove_pattern, w) is not None]

主函數：

if __name__ == '__main__':
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
	#讀本地文件
    data = spark.read.text('test')
    ＃使用廣播變量記錄總文本數量
    count = data.count()
    brocast_count = sc.broadcast(count)
	# show 一下數據
    data.show()
    # 根據空格進行文章編號和文本正文的切分，然後對文本正文調用jieba分詞進行分詞
    data = data.rdd.map(lambda x: x[0].split(' '))\
        .map(lambda x: (x[0], x[1], seg(x[1]))) # 分詞

過濾掉空行

如果我們的數據中存在空行或者分詞後去除停用詞後沒有詞了，那麼我們可以把這一行給去掉，spark中去掉空行可以使用filter，我習慣去掉空行的做法是使用一個特殊的符號，然後進行filter,代碼如下：

# 使用-- 標識空行，然後filter 掉這一行的數據，filter中如果改行的邏輯是false那麼就會filter 否則就pass
data = data\
    .map(lambda x: (x[0], '--' if len(x[2]) == 0 else x[1], x[2]))\
    .filter(lambda x: x[1] != '--')

計算TF

TF是針對一篇文章而言的，是一篇文章中得單詞頻率/單詞總數，這裏的計算比較簡單，就不多折騰了。

def calc_tf(line):
    """
    計算每個單詞在每篇文章的tf
    :param line:
    :return:
    """
    cnt_map = {}
    for w in line[2]:
        cnt_map[w] = cnt_map.get(w, 0) + 1
    lens = len(line)
    return [(line[0], (w, cnt *1.0/lens)) for w,cnt in cnt_map.items()]
# 計算tf
tf_data = data.flatMap(calc_tf)

打印可以看到：

print(tf_data.collect())

[('1', ('來到', 0.3333333333333333)), ('1', ('北京', 0.3333333333333333)), ('1', ('清華大學', 0.3333333333333333)), ('2', ('來到', 0.3333333333333333)), ('2', ('網易', 0.3333333333333333)), ('2', ('杭研', 0.3333333333333333)), ('2', ('大廈', 0.3333333333333333)), ('3', ('來到', 0.3333333333333333)), ('3', ('北京', 0.3333333333333333)), ('3', ('清華大學', 0.3333333333333333)), ('4', ('來到', 0.3333333333333333)), ('4', ('網易', 0.3333333333333333)), ('4', ('杭研', 0.3333333333333333)), ('4', ('大廈', 0.3333333333333333)), ('5', ('來到', 0.6666666666666666)), ('5', ('北京', 0.6666666666666666)), ('5', ('清華大學', 0.6666666666666666))]

計算IDF

IDF是逆文檔頻率，表示一個單詞出現在語料庫中出現的頻率，也就是一個單詞在多少篇文章中出現了。
下面就給出二個計算IDF的方法，在計算IDF的時候，flatMap需要帶上文章ID：

思路1. 分詞後的結果進行flatMap,轉化爲文章ID，單詞的RDD，然後進行計算，這裏順便了解一下combinedByKey

def flat_with_doc_id(data):
    """
    flat map的時候帶上文章ID
    :param data:
    :return:
    """
    return [(w, data[0]) for w in data[2]]
t = data.flatMap(flat_with_doc_id)\
    .combineByKey(lambda v: 1,
                  lambda x, v: x + 1,
                  lambda x, y: x+y)\
    .map(lambda x: (x[0], x[1]*1.0/brocast_count.value))

下面的代碼打印

print('t', t.collect())

可以看到

t [('來到', 1.2), ('北京', 0.8), ('清華大學', 0.8), ('網易', 0.4), ('杭研', 0.4), ('大廈', 0.4)]

如果沒毛病，計算DF是沒有錯的。

思路二.考慮到我們最後需要計算TF-IDF，如果第一步計算出TF後再結合這一步計算出來的IDF，那麼就不可避免的進行join操作，這個shuffle非常耗時耗力，我們應該儘量的避免，那麼可不可以通過RDD的transformer計算IDF呢，當然是可以的，下面提供一個我寫的，效率個人覺得還可以，在300W篇文本，單機16G內存在1小時內可以完成。

    def create_combiner(v):
        t = []
        t.append(v)
        return (t, 1)

    def merge(x, v):
        # x==>(list, count)
        t = []
        if x[0] is not None:
            t = x[0]
        t.append(v)
        return (t, x[1] + 1)

    def merge_combine(x, y):
        t1 = []
        t2 = []
        if x[0] is not None:
            t1 = x[0]
        if y[0] is not None:
            t2 = y[0]
        t1 = t1.extend(t2)
        return (t1, x[1] + y[1])

    def flat_map_2(line):
        rst = []
        idf_value = line[1][-1] * 1.0 / brocast_count.value
        for doc_pair in line[1][:-1]:
            print(doc_pair)
            for p in doc_pair:
                rst.append(Row(docId=p[0], token=line[0], tf_value=p[1], idf_value=idf_value, tf_idf_value=p[1] * idf_value))
        return rst

idf_rdd = tf_data.map(lambda x: (x[1][0], (x[0], x[1][1])))\
    .combineByKey(create_combiner,
                   merge,
                   merge_combine)\
    .flatMap(flat_map_2)

將其轉化爲DataFram然後show

    tf_idf_df = spark.createDataFrame(idf_rdd)
    tf_idf_df.show()
    tf_idf_df.printSchema()

可以看到

+-----+---------+-------------------+------------------+-----+
|docId|idf_value|       tf_idf_value|          tf_value|token|
+-----+---------+-------------------+------------------+-----+
|    1|      1.0| 0.3333333333333333|0.3333333333333333|   來到|
|    2|      1.0| 0.3333333333333333|0.3333333333333333|   來到|
|    3|      1.0| 0.3333333333333333|0.3333333333333333|   來到|
|    4|      1.0| 0.3333333333333333|0.3333333333333333|   來到|
|    5|      1.0| 0.6666666666666666|0.6666666666666666|   來到|
|    1|      0.6|0.19999999999999998|0.3333333333333333|   北京|
|    3|      0.6|0.19999999999999998|0.3333333333333333|   北京|
|    5|      0.6|0.39999999999999997|0.6666666666666666|   北京|
|    1|      0.6|0.19999999999999998|0.3333333333333333| 清華大學|
|    3|      0.6|0.19999999999999998|0.3333333333333333| 清華大學|
|    5|      0.6|0.39999999999999997|0.6666666666666666| 清華大學|
|    2|      0.4|0.13333333333333333|0.3333333333333333|   網易|
|    4|      0.4|0.13333333333333333|0.3333333333333333|   網易|
|    2|      0.4|0.13333333333333333|0.3333333333333333|   杭研|
|    4|      0.4|0.13333333333333333|0.3333333333333333|   杭研|
|    2|      0.4|0.13333333333333333|0.3333333333333333|   大廈|
|    4|      0.4|0.13333333333333333|0.3333333333333333|   大廈|
+-----+---------+-------------------+------------------+-----+

如果只是想計算DF，那麼可以直接使用第一個方法，combineByKey比reduceByKey要節省內存消耗，而且在大數據的時候更爲明顯。

完整代碼：

# -*- coding: utf-8 -*-

"""
 計算TF-IDF
 @Time    : 2019/2/18 18:03
 @Author  : MaCan ([email protected])
 @File    : text_transformator.py
"""

from pyspark.sql import SparkSession, Row
from pyspark.sql.types import *
from pyspark import SparkConf, SparkContext

# from spark_work.io_utils import mrp_hdfs_2016_path,mrp_hdfs_2017_path,mrp_hdfs_2018_path, user_dict_path

import jieba
import os
import re

#過濾英文的pattern
remove_pattern = '[\u4e00-\u9fa5]+'


# if os.path.exists(user_dict_path):
#     try:
#         jieba.load_userdict(user_dict_path)
#     except Exception as e:
#         print(e)


def seg(data):
    """
    分詞後返回分詞的dataframe
    :param spark:
    :param data:
    :return:
    """
    return [w for w in jieba.cut(data.strip(), cut_all=False) if len(w) > 1 and re.match(remove_pattern, w) is not None]


def flat_with_doc_id(data):
    """
    flat map的時候帶上文章ID
    :param data:
    :return:
    """
    return [(data[0], w) for w in data[2]]


def calc_tf(line):
    """
    計算每個單詞在每篇文章的tf
    :param line:
    :return:
    """
    cnt_map = {}
    for w in line[2]:
        cnt_map[w] = cnt_map.get(w, 0) + 1
    lens = len(line)
    return [(line[0], (w, cnt *1.0/lens)) for w,cnt in cnt_map.items()]

def create_combiner(v):
    t = []
    t.append(v)
    return (t, 1)

def merge(x, v):
    # x==>(list, count)
    t = []
    if x[0] is not None:
        t = x[0]
    t.append(v)
    return (t, x[1] + 1)

def merge_combine(x, y):
    t1 = []
    t2 = []
    if x[0] is not None:
        t1 = x[0]
    if y[0] is not None:
        t2 = y[0]
    t1 = t1.extend(t2)
    return (t1, x[1] + y[1])

def flat_map_2(line):
    rst = []
    idf_value = line[1][-1] * 1.0 / brocast_count.value
    for doc_pair in line[1][:-1]:
        print(doc_pair)
        for p in doc_pair:
            rst.append(Row(docId=p[0], token=line[0], tf_value=p[1], idf_value=idf_value, tf_idf_value=p[1] * idf_value))
    return rst


if __name__ == '__main__':
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()

    data = spark.read.text('test')
    count = data.count()
    brocast_count = sc.broadcast(count)

    data.show()
    data = data.rdd.map(lambda x: x[0].split(' '))\
        .map(lambda x: (x[0], x[1], seg(x[1]))) # 分詞

    # filter掉返回false的結果
    data = data\
        .map(lambda x: (x[0], '--' if len(x[2]) == 0 else x[1], x[2]))\
        .filter(lambda x: x[1] != '--')
    print('data', data.collect())
    data.cache() # x[0]=> docId x[1]==> token
    # 算tf
    tf_data = data.flatMap(calc_tf)
    print(tf_data.collect())


    #idf
    t = data.flatMap(flat_with_doc_id)\
        .map(lambda x: (x[1], x[0]))\
        .combineByKey(lambda v: 1,
                      lambda x, v: x + 1,
                      lambda x, y: x+y)\
        .map(lambda x: (x[0], x[1]*1.0/brocast_count.value)).collect()
    print('t', t)
    print('*'*20)

    idf_rdd = tf_data.map(lambda x: (x[1][0], (x[0], x[1][1])))\
        .combineByKey(create_combiner,
                       merge,
                       merge_combine)\
        .flatMap(flat_map_2)

    tf_idf_df = spark.createDataFrame(idf_rdd)
    tf_idf_df.show()
    tf_idf_df.printSchema()
    spark.stop()

PySpark TF-IDF計算（2）

使用PySpark進行TF-IDF計算

1. 準備數據

2. 加載數據並且分詞

過濾掉空行

計算TF

計算IDF

完整代碼：

Python實現大麥網搶票的四大關鍵技術點解析

salesforce零基礎學習（一百三十八）零碎知識點小總結（十）

Windows下PySpark 環境搭建篇以及詞頻統計（1）

解決TensorRT編譯時protobuf模塊編譯錯誤

windows下python安裝scipy庫的方法

IDEA遠程調試hadoop

python 使用uwsgi 開啓多進程服務

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結