Как получить подмножество строк из DataFrame на основе условия?

Question

У меня есть два кадра данных:

edges =
   srcId    dstId    timestamp
   1        4        1346564657
   1        2        1345769687
   2        4        1345769687
   4        1        1345769687

vertices =
   id   name   s_type
   1    abc    A
   2    def    B
   3    rtf    C
   4    wrr    D

Я хочу получить подмножество vertices, которое будет включать id, которые не упомянуты среди srcId и dstId в edges.

Это ожидаемый результат:

sub_vertices =
   id   name   s_type
   3    rtf    C

Как я могу это сделать?

val sub_vertices = vertices
  .join(edges, col("id") =!= col("srcId") && col("id") =!= col("dstId"), "left")
  .na.fill(0)
  .drop("srcId","dstId", "timestamp")
  .dropDuplicates()

Это мой текущий код, но он дает неверный результат.

Answer 1

Вы можете собрать srdId и dstIds в набор и транслировать набор в , используемый в фильтре вершин dataframe как

import org.apache.spark.sql.functions._
//collect all the srdId and dstId from edges dataframe into set and broadcast
val srdIdList = sc.broadcast(edges.select(collect_set("srcId").as("collectSrc"), collect_set("dstId").as("collectDst")).rdd.map(row => row.getAs[Seq[Int]](0) ++ row.getAs[Seq[Int]](1) toSet).collect()(0))

//using udf function remove all the rows that are in collected srdId and dstId in above step
def containsUdf = udf((id: Int) => !srdIdList.value.contains(id))
vertices.filter(containsUdf(col("id"))).show(false)

, который должен дать вам

+---+----+------+
|id |name|s_type|
+---+----+------+
|3  |rtf |C     |
+---+----+------+

Answer 2

Вы были почти рядом, вот несколько вещей, которые вам нужно изменить

val sub_vertices = vertices
  .join(edges, col("id") === col("srcId") || col("id") === col("dstId") , "left")
  .filter($"srcId".isNull && $"dstId".isNull)
  .drop("srcId","dstId", "timestamp")

Вывод:

+---+----+------+
|id |name|s_type|
+---+----+------+
|3  |rtf |C     |
+---+----+------+