pyspark 中的 ARRAY_CONTAINS 多个值

Question

我正在使用pyspark.sql.dataframe.DataFrame。我想根据多个变量过滤stack 的行，而不是单个变量{val}。我正在使用 Python 2 Jupyter 笔记本。目前，我执行以下操作：

stack = hiveContext.sql("""
    SELECT * 
    FROM db.table
    WHERE col_1 != ''
""")

stack.show()
+---+-------+-------+---------+
| id| col_1 | . . . | list    |
+---+-------+-------+---------+
| 1 |   524 | . . . |[1, 2]   |
| 2 |   765 | . . . |[2, 3]   |
.
.
.
| 9 |   765 | . . . |[4, 5, 8]|

for i in len(list):
    filtered_stack = stack.filter("array_contains(list, {val})".format(val=val.append(list[i])))
    (some query on filtered_stack)

如何在 Python 代码中重写它以根据多个值过滤行？即其中 {val} 等于某个包含一个或多个元素的数组。

我的问题与：ARRAY_CONTAINS muliple values in hive 相关，但是我试图在 Python 2 Jupyter 笔记本中实现上述目标。

Answer 1

没有 UDF

import pyspark.sql.functions as F

vals = {1, 2, 3}

_ = F.array_intersect(
    F.col("list"), 
    F.array([F.lit(i) for i in vals])
)
# This will now give a boolean field for any row with a list which has values in vals
_ = F.size(_) > 0

Answer 2

使用 Python UDF：

from pyspark.sql.functions import udf, size
from pyspark.sql.types import *

intersect = lambda type: (udf(
    lambda x, y: (
        list(set(x) & set(y)) if x is not None and y is not None else None),
    ArrayType(type)))

df = sc.parallelize([([1, 2, 3], [1, 2]), ([3, 4], [5, 6])]).toDF(["xs", "ys"])

integer_intersect = intersect(IntegerType())

df.select(
    integer_intersect("xs", "ys"),
    size(integer_intersect("xs", "ys"))).show()

+----------------+----------------------+
|<lambda>(xs, ys)|size(<lambda>(xs, ys))|
+----------------+----------------------+
|          [1, 2]|                     2|
|              []|                     0|
+----------------+----------------------+

用文字：

from pyspark.sql.functions import array, lit

df.select(integer_intersect("xs", array(lit(1), lit(5)))).show()

+-------------------------+
|<lambda>(xs, array(1, 5))|
+-------------------------+
|                      [1]|
|                       []|
+-------------------------+

或

df.where(size(integer_intersect("xs", array(lit(1), lit(5)))) > 0).show()

+---------+------+
|       xs|    ys|
+---------+------+
|[1, 2, 3]|[1, 2]|
+---------+------+