Question

在scala / spark代码中，我有1个Dataframe，其中包含一些行：

col1      col2
Abc       someValue1 
xyz       someValue2
lmn       someValue3
zmn       someValue4
pqr       someValue5
cda       someValue6

我的变量ArrayBuffer[String]包含[xyz,pqr,abc];

我想基于 col1 中arraybuffer中的给定值来过滤给定的数据帧。

在SQL中它就像：

select * from tableXyz where col1 in("xyz","pqr","abc");

Answer 1

假设您拥有数据框：

val df = sc.parallelize(Seq(("abc","someValue1"),
                            ("xyz","someValue2"),
                            ("lmn","someValue3"),
                            ("zmn","someValue4"),
                            ("pqr","someValue5"),
                            ("cda","someValue6")))
                 .toDF("col1","col2") 

+----+----------+
|col1|      col2|
+----+----------+
| abc|someValue1|
| xyz|someValue2|
| lmn|someValue3|
| zmn|someValue4|
| pqr|someValue5|
| cda|someValue6|
+----+----------+

然后，您可以定义UDF以根据数组的值过滤数据帧：

val array = ArrayBuffer[String]("xyz","pqr","abc") 

val function: (String => Boolean) = (arg: String) => array.contains(arg)
val udfFiltering = udf(function)

val filtered = df.filter(udfFiltering(col("col1")))
filtered.show()
+----+----------+
|col1|      col2|
+----+----------+
| abc|someValue1|
| xyz|someValue2|
| pqr|someValue5|
+----+----------+

或者，您可以通过SQLContext注册数据帧并对其进行sql查询：

var elements = ""
array.foreach { el => elements += "\"" + el + "\"" + "," }
elements = elements.dropRight(1)    
val query = "select * from tableXyz where col1 in(" + elements + ")"

df.registerTempTable("tableXyz")
val filtered = sqlContext.sql(query)
filtered.show()
+----+----------+
|col1|      col2|
+----+----------+
| abc|someValue1|
| xyz|someValue2|
| pqr|someValue5|
+----+----------+

如何在scala中基于列值（通过arraybuffer的多个值）过滤数据帧

1 个答案: