Jak filtrować ramkę danych Spark, jeśli jedna kolumna jest członkiem innej kolumny - scala, apache-spark, apache-spark-sql, spark-dataframe

Mam ramkę danych z dwiema kolumnami (jeden ciąg i jeden ciąg znaków):

root
|-- user: string (nullable = true)
|-- users: array (nullable = true)
|    |-- element: string (containsNull = true)

Jak mogę filtrować ramkę danych tak, aby wynikowa ramka danych zawierała tylko te wiersze user jest w users?

Odpowiedzi:

Oczywiście, jest to możliwe i nie tak trudne. Aby to osiągnąć, możesz użyć UDF.

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._

val df = sc.parallelize(Array(
("1", Array("1", "2", "3")),
("2", Array("1", "2", "2", "3")),
("3", Array("1", "2"))
)).toDF("user", "users")

val inArray = udf((id: String, array: scala.collection.mutable.WrappedArray[String]) => array.contains(id), BooleanType)

df.where(inArray($"user", $"users")).show()

Wynik to:

+----+------------+
|user|       users|
+----+------------+
|   1|   [1, 2, 3]|
|   2|[1, 2, 2, 3]|
+----+------------+

Szybki i prosty:

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

df.where(expr("array_contains(users, user)")