Filtering values that are not in a list using expr and filter

Viewed 44

I want to filter out rows in a dataframe where a column is not part of a list.

I am aware that I can use udf to go about this and it works.

def filterNegatives(val: Seq[String]): Seq[String] = {
    val.filter(v => !badList.contains(v))
}
val filterNegativesUdf = udf(filterNegatives _, ArrayType(StringType))

val cleanedDF = myDF.withColumn("pos" , filterNegativesUdf(col("allVals")))

Was wondering if there is a non udf way of achieving this.

I have tried the following and it works.

val cleanedDF = myDF.withColumn("pos", expr(s"filter(allVals, val -> val NOT IN ('badval1', 'badval2'))"))

but my list badList contains ~10 elements and I'd rather keep it the code clean by defining a list.

I have tried using the list inside filter in different variations, but all of them had some errors.

.withColumn("pos", expr(s"filter(allVals, val NOT IN ${badList}"))

//error:no viable alternative at input 'NOT IN List'     

Using - scala version 2.11

1 Answers

Consider using array_contains within your higher-order function filter as shown below.

val df = Seq(
  (1, Array("a", "b", "c", "d", "e")),
  (2, Array("h", "i", "j", "k")),
  (3, Array("u", "u", "v", "v", "w", "w"))
).toDF("id", "values")

val badList = Array("a", "e", "i", "o", "u")

On Spark 3.x:

df.
  withColumn("pos", filter($"values", v => !array_contains(lit(badList), v))).
  show
/*
+---+------------------+------------+
| id|            values|         pos|
+---+------------------+------------+
|  1|   [a, b, c, d, e]|   [b, c, d]|
|  2|      [h, i, j, k]|   [h, j, k]|
|  3|[u, u, v, v, w, w]|[v, v, w, w]|
+---+------------------+------------+
*/

On spark 2.4:

df.
  withColumn("bad_list", lit(badList)).
  withColumn("pos", expr("filter(values, v -> !array_contains(bad_list, v))")).
  drop("bad_list").
  show

Note that you could also consider using function array_except, but the catch is that any duplicates in the original array will be removed:

df.
  withColumn("pos", array_except($"values", lit(badList))).
  show
/*
+---+------------------+---------+
| id|            values|      pos|
+---+------------------+---------+
|  1|   [a, b, c, d, e]|[b, c, d]|
|  2|      [h, i, j, k]|[h, j, k]|
|  3|[u, u, v, v, w, w]|   [v, w]|
+---+------------------+---------+
*/
Related