Spark (Scala) How to access specifc row in dataframe by "key" and modify it

Viewed 53

I have two dataframes, one looks like this

+------------------------------------------------------------+
|docs                                                        |
+------------------------------------------------------------+
|{doc1.txt -> 1, doc2.txt -> 3, doc3.txt -> 5, doc4.txt -> 1}|
|{doc1.txt -> 2, doc2.txt -> 2, doc3.txt -> 4}               |
|{doc1.txt -> 3, doc2.txt -> 2, doc4.txt -> 2}               |
+------------------------------------------------------------+

and the other like

+--------------+----------+
|      Document|doc_length|
+--------------+----------+
|      doc1.txt|         0|
|      doc2.txt|         0|
|      doc3.txt|         0|
|      doc3.txt|         0|
|      doc4.txt|         0|
+-------------------------+

for sake of example the documents are in order, but in my use case I cannot expect them to be.

now I want to iterate through the first dataframe and update the values in the second as I go. I got a loop like this

df1.foreach(r =>
      for (keyValPair <- r(0).asInstanceOf[Map[String, Long]]) {
        // Something needs to happen here
      } )

In every iteration I want to take take the key of the key-value-pair to select a specific row in the second dataframe and then add the value to the doc_length, so my final output of df2.show() would look like EDIT: Later down the line I probably want to do other more complicated mathematical operations here then just summing all the values up, that's why I was trying to use the structure described above

+--------------+----------+
|      Document|doc_length|
+--------------+----------+
|      doc1.txt|         6|
|      doc2.txt|         7|
|      doc3.txt|         9|
|      doc4.txt|         0|
+-------------------------+

This doesn't look like it should be too hard, but I don't know how I can access specific rows of a dataframe, by using a specific column as a key, and change them

1 Answers

You can explode the map column and group by key to sum up the lengths:

val df2 = df.select(explode(col("val")))
    .groupBy(col("key").as("document"))
    .agg(sum("value").as("doc_length"))

df2.show
+--------+----------+
|document|doc_length|
+--------+----------+
|doc1.txt|         6|
|doc4.txt|         3|
|doc3.txt|         9|
|doc2.txt|         7|
+--------+----------+
Related