在多个数组类型元素之间筛选数据

Question

数据帧模式：

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元素，其中COL1元素为A，COL2元素为偶数

预期输出：[S]

---

我研究了各种函数。例如-array_position，但它只返回第一次出现。

有没有什么简单的方法，或者我必须分解数组？

Answer 1

假设您的条件适用于具有相同索引的数组元素，则自Spark 2.4.0以来，可以在SQL中使用lambda函数筛选数组，但这仍然没有通过其他语言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]|
+---+------+