CODEWITHSUNDEEP
scalaapache-sparkdataframeapache-spark-sql
Ronnie Phelps
Ronnie Phelps

Reputation: 133

Merging the results of a scala spark dataframe as an array of results in another dataframe's column

Is there a way to take the following two dataframes and join them by the col0 field producing the output below?

//dataframe1

val df1 = Seq(
  (1, 9, 100.1, 10),
).toDF("pk", "col0", "col1", "col2")

//dataframe2

val df2 = Seq(
  (1, 9 "a1", "b1"),
  (2, 9 "a2", "b2")
).toDF("pk", "col0", "str_col1", "str_col2")

//expected dataframe result

+---+-----+----+---------------------------+
| pk| col1|col2| new_arr_col               |
+---+-----+----+---------------------------+
|  1|100.1|  10|[[1,9,a1, b1],[2,9,a2, b2]]|
+---+-----+----+---------------------------+

Upvotes: 0

Views: 157

Answers (1)

Tzach Zohar
Tzach Zohar

Reputation: 37822

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

// creating new array column out of all df2 columns:
val df2AsArray = df2.select($"col0", array(df2.columns.map(col): _*) as "new_arr_col")

val result = df1.join(df2AsArray, "col0")
  .groupBy(df1.columns.map(col): _*) // grouping by all df1 columns
  .agg(collect_list("new_arr_col") as "new_arr_col") // collecting array of arrays
  .drop("col0")

result.show(false)
// +---+-----+----+--------------------------------------------------------+
// |pk |col1 |col2|new_arr_col                                             |
// +---+-----+----+--------------------------------------------------------+
// |1  |100.1|10  |[WrappedArray(2, 9, a2, b2), WrappedArray(1, 9, a1, b1)]|
// +---+-----+----+--------------------------------------------------------+

Upvotes: 1

Related Questions