scala - Spark DataFrame 中查找重叠数据的函数

Question

我有以下数据框:

我正在尝试查找一起乘坐过 3 趟以上火车的乘客。

所以对于上面的例子，ID为1和2的乘客乘坐同一趟车的次数超过3次[2,3,4,6]，ID为4和5的乘客乘坐同一趟车的次数超过3次三倍 [7,32,44,54]

有没有可以为此编写的 scala 函数？我尝试过相交类型函数，但我似乎无法将其应用于整个 DataFrame。

感谢您的帮助。

因此，对于预期的输出，我认为它将返回一个包含以下内容的 DataFrame:

我有一个大约有 15,000 行的 DataFrame

谢谢

Answer 1

您可以将自连接与 array_intersect 内置函数等结合起来:

import org.apache.spark.sql.functions.expr

val df = Seq(
  (1, Seq(1, 2, 3, 4, 6)),
  (2, Seq(2, 3, 4, 6, 7)),
  (3, Seq(1, 2, 5, 9, 100)),
  (4, Seq(11, 2, 4, 5, 7, 32, 44, 54)),
  (5, Seq(7, 12, 34, 32, 44, 54)),
  (6, Seq(5, 21)) 
).toDF("passengerId", "trainId")

df.as("d1").join(df.as("d2"), $"d1.passengerId" =!= $"d2.passengerId")
            .selectExpr("d1.passengerId as passengerId1", "d2.passengerId as passengerId2", "d1.trainId as trainId1", "d2.trainId as trainId2")
            .where("size(array_intersect(trainId1, trainId2)) > 3")
            .selectExpr("array_sort(array(passengerId1, passengerId2)) as ar")
            .distinct()
            .selectExpr("ar[0] as usr1", "ar[1] as usr2")
            .show()

// +----+----+
// |usr1|usr2|
// +----+----+
// |1   |2   |
// |4   |5   |
// +----+----+