Question

使用

from pyspark.sql import functions as f

和方法 f.agg 和 f.collect_set 我在 dataFrame 中创建了一个列colSet，如下所示：

+-------+--------+
| index | colSet |
+-------+--------+
|      1|[11, 13]|
|      2|  [3, 6]|
|      3|  [3, 7]|
|      4|  [2, 7]|
|      5|  [2, 6]|
+-------+--------+

现在，如何使用python /和pyspark来仅选择那些行，例如，3是colSet条目中数组的元素（通常情况下，只能有两个条目）！）？

我尝试过使用像这样的udf函数：

isInSet = f.udf( lambda vcol, val: val in vcol, BooleanType())

通过

进行调用

dataFrame.where(isInSet(f.col('colSet'), 3))

我也尝试从调用者中删除f.col并在isInSet的定义中使用它，但是没有用，我得到一个例外：

AnalysisException: cannot resolve '3' given input columns: [index, colSet]

对于如何在给定具有collect_set结果的行的情况下选择具有特定条目（甚至更好的子集!!!）的行，我们对此表示赞赏。

Answer 1

您的原始UDF很好，但要使用它，您需要将值3作为文字传递：

dataFrame.where(isInSet(f.col('colSet'), f.lit(3)))

但正如jxc在评论中指出的那样，使用array_contains可能是更好的选择：

dataFrame.where(f.array_contains(f.col('colSet'), 3))

我还没有做任何基准测试，但一般来说，使用PySpark中的UDF比使用内置函数要慢，因为JVM和Python解释器之间来回通信。

Answer 2

我今天找到了解决方案（在星期五晚上失败后），而没有使用udf方法：

[3 in x[0] for x in list(dataFrame.select(['colSet']).collect())]

希望将来可以帮助其他人。

使用pyspark在collect_set之后的set中按值选择行

2 个答案: