如何快速检查 PySpark 数据框中是否存在行？

Question

I have a PySpark dataframe like this:

+------+------+
|     A|     B|
+------+------+
|     1|     2|
|     1|     3|
|     2|     3|
|     2|     5|
+------+------+

我想在表格上进行查找，以查看特定行是否存在。例如，对于测试条件A=2，B=5，代码应返回True；对于测试条件A=2，B=10，代码应返回False。

我尝试了以下代码：

df[(df['A'] == 1) & (df['B'] == 2)].rdd.isEmpty()

很不幸，这段代码执行的时间很长，而且由于这是一次查找，会多次执行（针对不同的A和B值），我希望有一个更快的方法来完成这个任务。

我正在考虑的其他解决方案包括：

• 将PySpark数据框转换为Pandas数据框，因为行查找速度更快
• 使用.where()或.filter()，但从我尝试过的情况来看，我不预计它们会更快
• 使用.count()代替isEmpty()

- raysincanes

Answer 1

最好从你想要查找的条目中创建一个spark dataframe，然后执行一个半连接（semi join）或反连接（anti join）以获取已存在或不存在于查找dataframe中的行。这比逐个检查条目更有效率。

import pyspark.sql.functions as F

df = spark.createDataFrame([[2,5],[2,10]],['A','B'])

result1 = df.join(lookup, ['A','B'], 'semi').withColumn('exists', F.lit(True))

result2 = df.join(lookup, ['A','B'], 'anti').withColumn('exists', F.lit(False))

result = result1.unionAll(result2)

result.show()
+---+---+------+
|  A|  B|exists|
+---+---+------+
|  2|  5|  true|
|  2| 10| false|
+---+---+------+

- mck

Answer 2

Spark函数ANY提供了一种非常快速的方法来检查数据框中是否存在记录。

check = df.selectExpr('ANY((A = 2) AND (B = 5)) as chk')

check.show()
#  +----+
#  | chk|
#  +----+
#  |true|
#  +----+

check = df.selectExpr('ANY((A = 2) AND (B = 10)) as chk')
check.show()
#  +-----+
#  |  chk|
#  +-----+
#  |false|
#  +-----+

- ZygD