如何在PySpark中将两列堆叠成一列？

Question

我有以下PySpark DataFrame：

id   col1   col2
A    2      3
A    2      4
A    4      6
B    1      2

我想要将col1和col2叠加在一起，以便获得如下的单列:

id   col3
A    2   
A    3
A    4
A    6
B    1
B    2

我该如何做到呢？

df = (
    sc.parallelize([
        (A, 2, 3), (A, 2, 4), (A, 4, 6),
        (B, 1, 2),
    ]).toDF(["id", "col1", "col2"])
)

- Fluxy

---

你有漏掉什么吗？在col3中，你只有6个元素，而在原始的两个中总共有8个元素。 - Sergey Ronin

@KomronAripov：这些应该是唯一的值。 - Fluxy

Answer 1

最简单的方法是将col1和col2合并成一个数组列，然后使用explode函数进行拆分：

df.show()
+---+----+----+
| id|col1|col2|
+---+----+----+
|  A|   2|   3|
|  A|   2|   4|
|  A|   4|   6|
|  B|   1|   2|
+---+----+----+

df.selectExpr('id', 'explode(array(col1, col2))').show()
+---+---+
| id|col|
+---+---+
|  A|  2|
|  A|  3|
|  A|  2|
|  A|  4|
|  A|  4|
|  A|  6|
|  B|  1|
|  B|  2|
+---+---+

如果不需要重复项，您可以删除它们。

- Psidom

Answer 2

为了实现这个目的，需要按“id”分组，然后在聚合中收集“col1”和“col2”中的列表，再将其展开到一列中。 要获取唯一的数字，只需在之后删除重复项即可。

我看到您的最终结果中数字也已经排序，这是通过在聚合中对连接的列表进行排序来完成的。

以下是代码：

from pyspark.sql.functions import concat, collect_list, explode, col, sort_array

df = (
    sc.parallelize([
        ('A', 2, 3), ('A', 2, 4), ('A', 4, 6),
        ('B', 1, 2),
    ]).toDF(["id", "col1", "col2"])
)

result = df.groupBy("id") \
.agg(sort_array(concat(collect_list("col1"),collect_list("col2"))).alias("all_numbers")) \
.orderBy("id") \
.withColumn('number', explode(col('all_numbers'))) \
.dropDuplicates() \
.select("id","number") \
.show()

会产生如下结果：

+---+------+
| id|number|
+---+------+
|  A|     2|
|  A|     3|
|  A|     4|
|  A|     6|
|  B|     1|
|  B|     2|
+---+------+

- Sander Lam

Answer 3

如果涉及的列数较少，那么这是一个相对简单的解决方案。

df = (
    sc.parallelize([
        ('A', 2, 3), ('A', 2, 4), ('A', 4, 6),
        ('B', 1, 2),
    ]).toDF(["id", "col1", "col2"])
)


df.show()

+---+----+----+
| id|col1|col2|
+---+----+----+
|  A|   2|   3|
|  A|   2|   4|
|  A|   4|   6|
|  B|   1|   2|
+---+----+----+

df1 = df.select(['id', 'col1'])
df2 = df.select(['id', 'col2']).withColumnRenamed('col2', 'col1')

df_new = df1.union(df2)
df_new = df_new.drop_duplicates()
df_new.show()

+---+----+
| id|col1|
+---+----+
|  A|   3|
|  A|   4|
|  B|   1|
|  A|   6|
|  A|   2|
|  B|   2|
+---+----+

- Preetham