基于其中一个数组中的 Null 值共同过滤 Pyspark 结构中的两个数组

Question

我想在具有日期和值字段的结构中过滤两个有序数组。下面的示例 DataFrame 后面是解释和我正在尝试做的示例。

from pyspark.sql import Row
import datetime

rows = [
    Row(
        id ='1111',
        A=Row(
            dates=[datetime.datetime(2015, 7, 29, 14, 27), datetime.datetime(2015, 7, 31, 14, 27)],
            values=[20.0, 100.0]),
        B=Row(
            dates=[datetime.datetime(2015, 4, 18, 17, 52)],
            values=[12.58])
    ),
    Row(
        id='2222',
        A=Row(
            dates=[datetime.datetime(2011, 4, 28, 14, 27), datetime.datetime(2011, 4, 28, 14, 27)],
            values=[100.0, None]),
        B=Row(
            dates=[datetime.datetime(2011, 4, 18, 17, 52)],
            values=[None])
    ),
    Row(
        id='3333',
        A=None,
        B=None)
]

df = spark.createDataFrame(rows)

df.show(10, False)

+----+-----------------------------------------------------------+--------------------------------+
|id  |A                                                          |B                               |
+----+-----------------------------------------------------------+--------------------------------+
|1111|[[2015-07-29 14:27:00, 2015-07-31 14:27:00], [20.0, 100.0]]|[[2015-04-18 17:52:00], [12.58]]|
|2222|[[2011-04-28 14:27:00, 2011-04-28 14:27:00], [100.0,]]     |[[2011-04-18 17:52:00], []]     |
|3333|null                                                       |null                            |
+----+-----------------------------------------------------------+--------------------------------+

df.printSchema()

root
 |-- id: string (nullable = true)
 |-- A: struct (nullable = true)
 |    |-- dates: array (nullable = true)
 |    |    |-- element: timestamp (containsNull = true)
 |    |-- values: array (nullable = true)
 |    |    |-- element: double (containsNull = true)
 |-- B: struct (nullable = true)
 |    |-- dates: array (nullable = true)
 |    |    |-- element: timestamp (containsNull = true)
 |    |-- values: array (nullable = true)
 |    |    |-- element: double (containsNull = true)

这里对日期和值进行排序，以便日期的每个元素与值的每个元素对应。所以 values 中的第二个值与 dates 中的第二个日期一致。

我想过滤 df 以便它删除 Null 及其对应的日期，如果值只有 Null，它只返回 Null。喜欢，

new_df.show(10, False)

+----+-----------------------------------------------------------+--------------------------------+
|id  |A                                                          |B                               |
+----+-----------------------------------------------------------+--------------------------------+
|1111|[[2015-07-29 14:27:00, 2015-07-31 14:27:00], [20.0, 100.0]]|[[2015-04-18 17:52:00], [12.58]]|
|2222|[[2011-04-28 14:27:00], [100.0]]                           |null                          |
|3333|null                                                       |null                            |
+----+-----------------------------------------------------------+--------------------------------+

Answer 1

您可以使用arrays_zip 和filter 进行空值过滤，然后将数组解压缩回其原始状态，同时处理空数组或空列：

import pyspark.sql.functions as F

df2 = df.withColumn(
    'A', 
    F.expr("filter(arrays_zip(A.dates, A.values), x -> x.values is not null)")
).withColumn(
    'A', 
    F.when(
        (F.size('A') != 0) & (F.col('A').isNotNull()), 
        F.struct(F.col('A.0').alias('dates'), F.col('A.1').alias('values'))
    )
).withColumn(
    'B', 
    F.expr("filter(arrays_zip(B.dates, B.values), x -> x.values is not null)")
).withColumn(
    'B', 
    F.when(
        (F.size('B') != 0) & (F.col('B').isNotNull()), 
        F.struct(F.col('B.0').alias('dates'), F.col('B.1').alias('values'))
    )
)

df2.show(truncate=False)
+----+-----------------------------------------------------------+--------------------------------+
|id  |A                                                          |B                               |
+----+-----------------------------------------------------------+--------------------------------+
|1111|[[2015-07-29 13:27:00, 2015-07-31 13:27:00], [20.0, 100.0]]|[[2015-04-18 16:52:00], [12.58]]|
|2222|[[2011-04-28 13:27:00], [100.0]]                           |null                            |
|3333|null                                                       |null                            |
+----+-----------------------------------------------------------+--------------------------------+

df2.printSchema()
root
 |-- id: string (nullable = true)
 |-- A: struct (nullable = true)
 |    |-- dates: array (nullable = true)
 |    |    |-- element: timestamp (containsNull = true)
 |    |-- values: array (nullable = true)
 |    |    |-- element: double (containsNull = true)
 |-- B: struct (nullable = true)
 |    |-- dates: array (nullable = true)
 |    |    |-- element: timestamp (containsNull = true)
 |    |-- values: array (nullable = true)
 |    |    |-- element: double (containsNull = true)