如果一列是另一列的成员,如何过滤Spark数据帧

时间:2016-09-10 03:09:43

标签: scala apache-spark dataframe apache-spark-sql

我有一个包含两列的数据框(一个字符串和一个字符串数组):

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

如何过滤数据框,以便结果数据框只包含user位于users的行?

2 个答案:

答案 0 :(得分:10)

快速而简单:

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

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

答案 1 :(得分:5)

当然,这是可能的,而不是那么难。为此,您可以使用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()

输出结果为:

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