在多个数组类型元素中过滤数据

时间:2019-11-01 11:38:20

标签: apache-spark pyspark

数据框架构:

root
 |-- ID: decimal(15,0) (nullable = true)
 |-- COL1: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- COL2: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- COL3: array (nullable = true)
 |    |-- element: string (containsNull = true)

样本数据

+--------------------+--------------------+--------------------+
|         COL1       |        COL2        |              COL3  |
+--------------------+--------------------+--------------------+
|[A, B, C, A]        |[101, 102, 103, 104]|[P, Q, R, S]        |
+--------------------+--------------------+--------------------+

我想在数组元素上应用嵌套条件。

例如

查找COL3个元素为{A和COL1个元素为偶数的COL2个元素。

预期输出:[S]


我研究了各种功能。例如-array_position,但它仅返回第一次出现的情况。

有什么简单的方法还是必须爆炸数组?

1 个答案:

答案 0 :(得分:2)

假设条件适用于具有相同索引的数组元素,则自Spark 2.4.0起可以在SQL中使用lambda函数filter个数组,但这仍未通过其他语言API公开,您需要使用expr()。您只需压缩三个数组,然后过滤生成的结构体数组:

scala> df.show()
+---+------------+--------------------+------------+
| ID|        COL1|                COL2|        COL3|
+---+------------+--------------------+------------+
|  1|[A, B, C, A]|[101, 102, 103, 104]|[P, Q, R, S]|
+---+------------+--------------------+------------+

scala> df.select($"ID", expr(s"""
     | filter(
     |   arrays_zip(COL1, COL2, COL3),
     |   e -> e.COL1 == "A" AND CAST(e.COL2 AS integer) % 2 == 0
     | ).COL3 AS result
     | """)).show()
+---+------+
| ID|result|
+---+------+
|  1|   [S]|
+---+------+

由于它使用expr()来提供SQL表达式作为列,因此它也可以与PySpark一起使用:

>>> from pyspark.sql.functions import expr
>>> df.select(df.ID, expr("""
...   filter(
...     arrays_zip(COL1, COL2, COL3),
...     e -> e.COL1 == "A" AND CAST(e.COL2 AS integer) % 2 == 0
...   ).COL3 AS result
... """)).show()
+---+------+
| ID|result|
+---+------+
|  1|   [S]|
+---+------+