Spark Scala UDF to count number of array elements contained in another string column - scala

Question

I have a spark dataframe df with 2 columns, say A and B, where A is array of string type and B is a string.
For each row, I am trying to count how many elements in A are contained in B. The UDF I have written is as follows. I thought it should be easy but it breaks down in the subsequent action step.
val hasAddressInUDF = udf{(s: String, t: Array[String]) => t.filter(word => s.contains(word)).size}
Could anyone help? Thanks.

Answer 1

You should be able to use the Seq.count method for this one: for each element in the sequence you can count whether it is a substring of your B column.
import spark.implicits._
val df = Seq(
(Seq("potato", "meat", "car"), "potatoes with meat"),
(Seq("ice", "cream", "candy"), "pasta with cream"),
(Seq("crackers", "with", "hummus"), "tasty food")
).toDF("A", "B")
df.show(false)
+------------------------+------------------+
|A |B |
+------------------------+------------------+
|[potato, meat, car] |potatoes with meat|
|[ice, cream, candy] |pasta with cream |
|[crackers, with, hummus]|tasty food |
+------------------------+------------------+
def count_in_seq = udf((A: Seq[String], B: String) => A.count(elem => B.contains(elem)))
df.withColumn("counts", count_in_seq($"A", $"B")).show(false)
+------------------------+------------------+------+
|A |B |counts|
+------------------------+------------------+------+
|[potato, meat, car] |potatoes with meat|2 |
|[ice, cream, candy] |pasta with cream |1 |
|[crackers, with, hummus]|tasty food |0 |
+------------------------+------------------+------+
Hope this helps!