scala – 按某些列值拆分Spark数据帧,然后独立于其他列旋转每个

val inputDF = Seq(("tom","20","a","street a","germany"),("jimmy","30","b","street b",("lola","50","c","street c","argentina"),("maria","60","d","street d",("joe","70","e","street e","argentina")
                  .toDF("name","age","company","address","country")

//+-----+---+-------+--------+---------+
//| name|age|company| address|  country|
//+-----+---+-------+--------+---------+
//|  tom| 20|      a|street a|  germany|
//|jimmy| 30|      b|street b|  germany|
//| lola| 40|      c|street c|argentina|
//|maria| 50|      d|street d|argentina|
//|  joe| 60|      e|street e|argentina|
//+-----+---+-------+--------+---------+

我需要按“国家/地区”列的不同值拆分记录.对于输入数据帧,拆分应该产生：

//+-----+---+-------+--------+---------+
//| name|age|company| address|  country|
//+-----+---+-------+--------+---------+
//|  tom| 20|      a|street a|  germany|
//|jimmy| 30|      b|street b|  germany|
//+-----+---+-------+--------+---------+

//+-----+---+-------+--------+---------+
//| name|age|company| address|  country|
//+-----+---+-------+--------+---------+
//| lola| 40|      c|street c|argentina|
//|maria| 50|      d|street d|argentina|
//|  joe| 60|      e|street e|argentina|
//+-----+---+-------+--------+---------+

而且我还必须在每个数据框下旋转“名称”和“年龄”列,以便每个人拥有不同的公司和地址,同时保持其余列的完整性.所需的输出数据框如下所示：

//+-----+---+-------+--------+---------+
//| name|age|company| address|  country|
//+-----+---+-------+--------+---------+
//|jimmy| 30|      a|street a|  germany|
//|  tom| 20|      b|street b|  germany|
//|  joe| 60|      c|street c|argentina|
//| lola| 40|      d|street d|argentina|
//|maria| 50|      e|street e|argentina|
//+-----+---+-------+--------+---------+

最后的行顺序无关紧要

我的第一次尝试(在Spark-shell上运行)

我尝试为每一行分配一个唯一的id,然后将所需的列(名称和年龄)洗牌,并使用辅助id值将重新排序的数据帧与其余的数据帧连接起来.这里的主要问题是使用collect(),这对于大数据帧可能是危险的,并且重新分区(1)几乎违反分布式计算和Spark(它用于在使用不同数量的分区压缩rdds时避免异常) .

import org.apache.spark.sql.{DataFrame,Row}
import org.apache.spark.sql.functions.{col,monotonically_increasing_id,rand}
import org.apache.spark.sql.types.LongType

// column(s) names to split the input dataframe
val colToSplit = Seq("country")
val splitCols = colToSplit.map(col)

// list of columns names to be rotated (together)
val colsToRotate = Seq("name","age")
val rotateCols = colsToRotate.map(col) :+ col(auxCol)

// add an auxiliar column for joining the dataframe in the final step
val auxCol = "aux"
val dfWithID = inputDF.withColumn(auxCol,monotonically_increasing_id())

val splitValuesSchema = dfWithID.select(splitCols: _*).schema

// create one dataframe for each value of the splitting column
val splitValuesDFs = dfWithID.select(splitCols: _*).distinct().collect()
  .map(row => spark.sparkContext.makeRDD(List(row)))
  .map(rdd => spark.createDataFrame(rdd,splitValuesSchema))

val rotateIDCols = Array(auxCol) ++ colsToRotate

// join the split values with their records (DFs with id + colsToRotate)
val splittedDFs = splitValuesDFs
  .map(df => df.join(dfWithID,colToSplit).selectExpr(rotateIDCols: _*))

// random reorder the auxiliar id column (DFs with random ids)
val randIdDFs = splittedDFs
  .map(df => df.select(auxCol).orderBy(rand()).toDF())

// get rdds with random ids
val randIdRdds = randIdDFs
  .map(df => df.select(auxCol).rdd.map(row => row(0)))

// finally,zip and append the rdds with the random ids to the dataframes created by
// splitting the main df to obtain the rotated dataframe with all the data
val tuples = (splittedDFs,randIdRdds).zipped
val newRdds = tuples
  .map((df: DataFrame,rdd) => df.rdd.repartition(1).zip(rdd.repartition(1))
  .map(row => Row.fromSeq(row._1.toSeq ++ Seq(row._2))))

val tuples2 = (splittedDFs,newRdds).zipped
val rotatedDF = tuples2.map((df: DataFrame,rdd) => spark
  .createDataFrame(rdd,df.schema.add("rotated_id",LongType)).drop(auxCol))
  .reduce(_ union _).withColumnRenamed("rotated_id","column2join")

// get the rest of the columns
val noRotateCols = dfWithID.columns.diff(colsToRotate).map(col)
val noRotatedDF = dfWithID.select(noRotateCols: _*)
  .withColumnRenamed(auxCol,"column2join")

// join both dataframes
val outputDF = noRotatedDF.join(rotatedDF,"column2join")
  .select(inputDF.columns.map(col): _*) // to keep the initial columns order

显示输出数据帧会产生与上述预期输出类似的结果(它主要取决于rand()函数顺序)

我想尽可能避免使用收集和重新分区,以获得更实用的解决方案.

欢迎任何评论或想法！