PySpark：计算列表中不同元素的出现次数

Question

我有以下数据：

data = {'date': ['2014-01-01', '2014-01-02', '2014-01-03', '2014-01-04', '2014-01-05', '2014-01-06'],
     'flat': ['A;A;B', 'D;P;E;P;P', 'H;X', 'P;Q;G', 'S;T;U', 'G;C;G']}

data['date'] = pd.to_datetime(data['date'])

data = pd.DataFrame(data)
data['date'] = pd.to_datetime(data['date'])
spark = SparkSession.builder \
    .master('local[*]') \
    .config("spark.driver.memory", "500g") \
    .appName('my-pandasToSparkDF-app') \
    .getOrCreate()
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
spark.sparkContext.setLogLevel("OFF")

df=spark.createDataFrame(data)
new_frame = df.withColumn("list", F.split("flat", "\;"))

我想新增一列，用于存储每个不同元素的出现次数（按升序排序），另一列则存储最大值。

+-------------------+-----------+---------------------+-----------+----+
|               date| flat      | list                |occurrences|max |
+-------------------+-----------+---------------------+-----------+----+
|2014-01-01 00:00:00|A;A;B      |['A','A','B']        |[1,2]      |2   |
|2014-01-02 00:00:00|D;P;E;P;P  |['D','P','E','P','P']|[1,1,3]    |3   |
|2014-01-03 00:00:00|H;X        |['H','X']            |[1,1]      |1   |
|2014-01-04 00:00:00|P;Q;G      |['P','Q','G']        |[1,1,1]    |1   |
|2014-01-05 00:00:00|S;T;U      |['S','T','U']        |[1,1,1]    |1   |
|2014-01-06 00:00:00|G;C;G      |['G','C','G']        |[1,2]      |2   |  
+-------------------+-----------+---------------------+-----------+----+

非常感谢！

- Dominik

---

对你来说，列的顺序：出现次数重要吗？ - Sreeram TP

Answer 1

对于Spark2.4+，可以通过使用高阶函数transform和aggregate的一个表达式来实现此目标，而不需要多个groupBys和aggregations(因为它们在大数据中是昂贵的shuffle操作)。这应该是Spark2.4的规范解决方案。请注意保留html标记。

from pyspark.sql import functions as F
df=spark.createDataFrame(data)
df.withColumn("list", F.split("flat","\;"))\
  .withColumn("occurances", F.expr("""array_sort(transform(array_distinct(list), x-> aggregate(list, 0,(acc,t)->acc+IF(t=x,1,0))))"""))\
  .withColumn("max", F.array_max("occurances"))\
  .show()
+-------------------+---------+---------------+----------+---+
|               date|     flat|           list|occurances|max|
+-------------------+---------+---------------+----------+---+
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|    [1, 2]|  2|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]| [1, 1, 3]|  3|
|2014-01-03 00:00:00|      H;X|         [H, X]|    [1, 1]|  1|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]| [1, 1, 1]|  1|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]| [1, 1, 1]|  1|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|    [1, 2]|  2|
+-------------------+---------+---------------+----------+---+

- murtihash

Answer 2

你可以通过几个groupBy语句来实现这一点，

首先，你有一个像这样的数据框，

+-------------------+---------+---------------+
|               date|     flat|           list|
+-------------------+---------+---------------+
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|
|2014-01-03 00:00:00|      H;X|         [H, X]|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|
+-------------------+---------+---------------+

使用F.explode将list列拆分，如下所示：

new_frame_exp = new_frame.withColumn("exp", F.explode('list'))

然后，你的数据框将会是这样的：

+-------------------+---------+---------------+---+
|               date|     flat|           list|exp|
+-------------------+---------+---------------+---+
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|  A|
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|  A|
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|  B|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  D|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  P|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  E|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  P|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  P|
|2014-01-03 00:00:00|      H;X|         [H, X]|  H|
|2014-01-03 00:00:00|      H;X|         [H, X]|  X|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  P|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  Q|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  G|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  S|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  T|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  U|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|  G|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|  C|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|  G|
+-------------------+---------+---------------+---+

在这个数据框上，按照以下方式进行groupBy：

new_frame_exp_agg = new_frame_exp.groupBy('date', 'flat', 'list', 'exp').count()

然后你会得到一个像这样的数据框，

+-------------------+---------+---------------+---+-----+
|               date|     flat|           list|exp|count|
+-------------------+---------+---------------+---+-----+
|2014-01-03 00:00:00|      H;X|         [H, X]|  H|    1|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  G|    1|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  U|    1|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  T|    1|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  P|    1|
|2014-01-03 00:00:00|      H;X|         [H, X]|  X|    1|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|  G|    2|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  E|    1|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|  C|    1|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]|  S|    1|
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|  B|    1|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  D|    1|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]|  Q|    1|
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|  A|    2|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]|  P|    3|
+-------------------+---------+---------------+---+-----+

在这个数据框上，应用一层聚合来收集计数并列出最大值，如下所示：

res = new_frame_exp_agg.groupBy('date', 'flat', 'list').agg(
                                         F.collect_list('count').alias('occurances'),
                                         F.max('count').alias('max'))

res.orderBy('date').show()


+-------------------+---------+---------------+----------+---+
|               date|     flat|           list|occurances|max|
+-------------------+---------+---------------+----------+---+
|2014-01-01 00:00:00|    A;A;B|      [A, A, B]|    [2, 1]|  2|
|2014-01-02 00:00:00|D;P;E;P;P|[D, P, E, P, P]| [1, 1, 3]|  3|
|2014-01-03 00:00:00|      H;X|         [H, X]|    [1, 1]|  1|
|2014-01-04 00:00:00|    P;Q;G|      [P, Q, G]| [1, 1, 1]|  1|
|2014-01-05 00:00:00|    S;T;U|      [S, T, U]| [1, 1, 1]|  1|
|2014-01-06 00:00:00|    G;C;G|      [G, C, G]|    [1, 2]|  2|
+-------------------+---------+---------------+----------+---+

如果您想对列occurance进行排序，您可以在Spark 2.4+上使用F.array_sort函数对该列进行排序，否则您需要编写一个UDF来实现。

- Sreeram TP