数据框架构:
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,但它仅返回第一次出现的情况。
有什么简单的方法还是必须爆炸数组?
答案 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]|
+---+------+