Spark (Scala) Turn a list with duplicates into a map of (list_entry, count)

user_12943580 0
Viewed 245

So, my current code looks something along the lines of

val tokens =
    .read.option("wholetext", true).textFile(path)
    .flatMap(line => line.split(" ")
    .withColumn("filename", substring_index(input_file_name, "/", -1))
    .groupBy(col("value"))
    .agg(collect_list("filename") as "docs")

It reads in all text documents in path and creates a Dataframe that looks something like:

|word1   |[1.txt]                            |
|word2   |[2.txt]                            |
|word3   |[2.txt, 1.txt, 1.txt]              |

Now I would like to take the lists and reduce them to something like this (sorry I don't actually know what a map looks like in Scala)

|word1   |[(1.txt, 1)]                       |
|word2   |[(2.txt, 1)]                       |
|word3   |[(1.txt, 2), (2.txt, 1)]           |

In theory I know what to do: take the list entries, map them to (entry, 1) and then reduce them by adding up the count. But I have very little experience with Scala so I don't actually know how to write this into code.

As mentioned above, I would like to have the document names be the keys in a map, to make accessing the counts easier.

apache-spark scala
1 Answers

This should do it, get an Array, you can convert to Map. Here you get an Array of String, you would need to split to 2 fields based on comma and look for guidance at Spark Scala Dataframe convert a column of Array of Struct to a column of Map, which I upvoted. You can also use struct instead of the Array.

import org.apache.spark.sql.functions._
val df = spark.read.textFile("/FileStore/tables/fff1.txt", "/FileStore/tables/fff2.txt")
val df2 = df.flatMap(_.split(" ")).withColumn("filename", input_file_name).groupBy("Value", "filename").count()
val df3 = df2.groupBy("Value").agg(collect_list(array("filename", "count")) as "docs")
df3.show(false)

returns:

+-----+----------------------------------------------------------------------------+
|Value|docs                                                                        |
+-----+----------------------------------------------------------------------------+
|you  |[[dbfs:/FileStore/tables/fff2.txt, 2], [dbfs:/FileStore/tables/fff1.txt, 1]]|
|fine |[[dbfs:/FileStore/tables/fff2.txt, 1], [dbfs:/FileStore/tables/fff1.txt, 1]]|
|how  |[[dbfs:/FileStore/tables/fff2.txt, 1], [dbfs:/FileStore/tables/fff1.txt, 1]]|
|hear |[[dbfs:/FileStore/tables/fff2.txt, 1], [dbfs:/FileStore/tables/fff1.txt, 1]]|
|ok   |[[dbfs:/FileStore/tables/fff2.txt, 1]]                                      |
|have |[[dbfs:/FileStore/tables/fff2.txt, 1]]                                      |
...

For completeness, to get the Map:

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

val arrayToMap = udf[Map[String, Long], Seq[Row]] {
  array => array.map { case Row(key: String, value: Long) => (key, value) }.toMap
}

val df = spark.read.textFile("/FileStore/tables/fff1.txt", "/FileStore/tables/fff2.txt")
val df2 = df.flatMap(_.split(" ")).withColumn("filename", input_file_name).groupBy("Value", "filename").count()
val df3 = df2.groupBy("Value").agg(collect_list(struct("filename", "count")))   

val df4 = df3.withColumn("words", arrayToMap($"collect_list(struct(filename, count))"))
df4.show(false)
user_69339930
Related