有条件合并阵列

Question

我有这样的数据框：

column_1   column_2    column_3
   [1,3]        [2]           2
 [1,2,3]       null           1 
   [3,4]        [6]           1

如果不是null和del column_2，我想将列_2的值附加到column_1。

所需的输出：

column_1    column_3
 [1,3,2]           2
 [1,2,3]           1
 [3,4,6]           1

Answer 1

• 如果列_2为null，请返回column_1
• 否则，使用column_2使用array_union union column_1

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [([1,3], [2], 2),
     ([1,2,3], None, 1),
     ([3,4], [6], 1)],
    ['column_1', 'column_2', 'column_3']
)
df = df.select(
    F.when(F.isnull('column_2'), F.col('column_1')).otherwise(F.array_union('column_1', 'column_2')).alias('column_1'),
    'column_3'
)
df.show()
# +---------+--------+
# | column_1|column_3|
# +---------+--------+
# |[1, 3, 2]|       2|
# |[1, 2, 3]|       1|
# |[3, 4, 6]|       1|
# +---------+--------+