在PySpark中为每一行计算列中不同子字符串出现的次数？

Question

我的数据集看起来像这样，其中我有一个逗号分隔的字符串值集合在col1和col2中。col3是将这两列连接在一起的结果。

+===========+========+=============
|col1       |col2    |col3   
+===========+========+=============
|a,b,c,d    |a,c,d   |a,b,c,d,a,c,d 
|e,f,g      |f,g,h   |e,f,g,f,g,h
+===========+========+=============

基本上，我想做的是获取以逗号分隔的所有值，并在每个值及其计数后附加另一列。

我希望在col4中获得这种输出：

+===========+========+==============+======================
|col1       |col2    |col3          |col4
+===========+========+==============+======================
|a,b,c,d    |a,c,d   |a,b,c,d,a,c,d |a: 2, b: 1, c: 2, d: 2
|e,f,g      |f,g,h   |e,f,g,f,g,h   |e: 1, f: 2, g: 2, h: 1
+===========+========+==============+======================

我已经弄清楚了如何将列连接起来以得到col3，但是在得到col4时遇到了一些困难。这是我停下来的地方，对接下来该怎么做有点不确定：

from pyspark.sql.functions import concat, countDistinct

df = df.select(concat(df.col1, df.col2).alias('col3'), '*')
df.agg(countDistinct('col3')).show()
# +--------------------+ 
# |count(DISTINCT col3)|
# +--------------------+
# |                   2|
# +--------------------+

如何动态计算以逗号分隔的col3中的子字符串，并在数据集的所有行中创建一个显示每个子字符串频率的最终列？

- user16315820

Answer 1

使用UDF

以下是使用UDF的方法。 首先进行数据生成。

from pyspark.sql.types import StringType, StructType, StructField
from pyspark.sql.functions import concat_ws, udf

data = [("a,b,c,d", "a,c,d", "a,b,c,d,a,c,d"),
        ("e,f,g", "f,g,h", "e,f,g,f,g,h")
       ]
schema = StructType([
    StructField("col1",StringType(),True),
    StructField("col2",StringType(),True),
    StructField("col3",StringType(),True),
])
df = spark.createDataFrame(data=data,schema=schema)

然后使用一些本地的Python函数，如Counter和json来完成任务。

from collections import Counter
import json

@udf(StringType())
def count_occurances(row):
    return json.dumps(dict(Counter(row.split(','))))

df.withColumn('concat', concat_ws(',', df.col1, df.col2, df.col3))\
  .withColumn('counts', count_occurances('concat')).show(2, False)

结果为

+-------+-----+-------------+---------------------------+--------------------------------+
|col1   |col2 |col3         |concat                     |counts                          |
+-------+-----+-------------+---------------------------+--------------------------------+
|a,b,c,d|a,c,d|a,b,c,d,a,c,d|a,b,c,d,a,c,d,a,b,c,d,a,c,d|{"a": 4, "b": 2, "c": 4, "d": 4}|
|e,f,g  |f,g,h|e,f,g,f,g,h  |e,f,g,f,g,h,e,f,g,f,g,h    |{"e": 2, "f": 4, "g": 4, "h": 2}|
+-------+-----+-------------+---------------------------+--------------------------------+

使用原生pyspark函数的解决方案

这种解决方案比使用udf更复杂，但由于缺少udf，可能更具性能。思路是将三个字符串列连接起来并展开它们。为了知道每个展开行来自哪里，我们添加了一个索引。双重分组将帮助我们获得所需的结果。最后，我们将结果与原始框架连接以获得所需的模式。

from pyspark.sql.functions import concat_ws, monotonically_increasing_id, split, explode, collect_list
df = df.withColumn('index', monotonically_increasing_id())

df.join(
    df
  .withColumn('concat', concat_ws(',', df.col1, df.col2, df.col3))\
  .withColumn('arr_col', split('concat', ','))\
  .withColumn('explode_col', explode('arr_col'))\
  .groupBy('index', 'explode_col').count()\
  .withColumn('concat_counts', concat_ws(':', 'explode_col', 'count'))\
  .groupBy('index').agg(concat_ws(',', collect_list('concat_counts')).alias('grouped_counts')), on='index').show()

导致

+-----------+-------+-----+-------------+---------------+
|      index|   col1| col2|         col3| grouped_counts|
+-----------+-------+-----+-------------+---------------+
|42949672960|a,b,c,d|a,c,d|a,b,c,d,a,c,d|a:4,b:2,c:4,d:4|
|94489280512|  e,f,g|f,g,h|  e,f,g,f,g,h|h:2,g:4,f:4,e:2|
+-----------+-------+-----+-------------+---------------+

请注意，在UDF部分创建的JSON通常比使用本地pyspark函数的grouped_counts列中的简单字符串更方便使用。

- pythonic833

Answer 2

我建议使用原生的Spark选项，使用数组而不是字符串。

from pyspark.sql import functions as F
df = spark.createDataFrame([("a,b,c,d", "a,c,d"), ("e,f,g", "f,g,h")], ['col1', 'col2'])

conc = F.concat(F.split('col1', ','), F.split('col2', ','))
dist = F.array_distinct(conc)
m = F.map_from_arrays(dist, F.transform(dist, lambda x: F.size(F.filter(conc, lambda y: y == x))))
df = df.withColumn('counts', m)

df.show(truncate=0)
# +-------+-----+--------------------------------+
# |col1   |col2 |counts                          |
# +-------+-----+--------------------------------+
# |a,b,c,d|a,c,d|{a -> 2, b -> 1, c -> 2, d -> 2}|
# |e,f,g  |f,g,h|{e -> 1, f -> 2, g -> 2, h -> 1}|
# +-------+-----+--------------------------------+

输出的类型是 map，而不是 string。这样可以轻松通过例如 F.col('counts')['g'] 来访问值。

- ZygD