如何在PySpark数据帧的列（带有数据类型数组（字符串））上应用筛选器？

Question

我有一个PySpark数据帧：

df = spark.createDataFrame([
    ("u1", ['a', 'b']),
    ("u2", ['c', 'b']),
    ("u3", ['a', 'b']),

    ],
    ['user_id', 'features'])

print(df.printSchema())
df.show(truncate=False)

输出：

root
 |-- user_id: string (nullable = true)
 |-- features: array (nullable = true)
 |    |-- element: string (containsNull = true)

None
+-------+--------+
|user_id|features|
+-------+--------+
|u1     |[a, b]  |
|u2     |[c, b]  |
|u3     |[a, b]  |
+-------+--------+

我只想保留名为features[a，b]的列。由于此列是字符串数组，因此不能使用简单筛选器

我怎样才能做到这一点

预期产出：

+-------+--------+
|user_id|features|
+-------+--------+
|u1     |[a, b]  |
|u3     |[a, b]  |
+-------+--------+

Answer 1

您可以使用array(lit(...))

import pyspark.sql.functions as F

df2 = df.filter(F.array_sort(F.col('features')) == F.array_sort(F.array(F.lit('a'), F.lit('b'))))

df2.show()
+   -+    +
|user_id|features|
+   -+    +
|     u1|  [a, b]|
|     u3|  [a, b]|
+   -+    +