title: Split Spark dataframe and calculate average based on one column value
body: <p>I have two dataframes, the first dataframe <code>classRecord</code> has 10 different entries like the following:   </p>
  
  <pre><code>Class, Calculation
  first, Average
  Second, Sum
  Third, Average
  </code></pre>
  
  <p>Second dataframe <code>studentRecord</code> has around 50K entries like the following:   </p>
  
  <pre><code>Name, height, Camp, Class
  Shae, 152, yellow, first
  Joe, 140, yellow, first
  Mike, 149, white, first
  Anne, 142, red, first
  Tim, 154, red, Second
  Jake, 153, white, Second
  Sherley, 153, white, Second
  </code></pre>
  
  <p>From second dataframe, based on class type, I would like to perform calculation on height (for class first: average, for class second: sum, etc.) based on the camp separately (if class is first, avg of yellow, white and so on separately).
  I tried the following code:</p>
  
  <pre><code>//function to calculate average
  def averageOnName(splitFrame : org.apache.spark.sql.DataFrame ) : Array[(String, Double)] = {
    val pairedRDD: RDD[(String, Double)] = splitFrame.select($"Name",$"height".cast("double")).as[(String, Double)].rdd
    var avg_by_key = pairedRDD.mapValues(x =&gt; (x, 1)).reduceByKey((x, y) =&gt; (x._1 + y._1, x._2 + y._2)).mapValues(y =&gt; 1.0 * y._1 / y._2).collect
    return avg_by_key
  }
  
  //required schema for further modifications
  val schema = StructType(
  StructField("name", StringType, false) ::
  StructField("avg", DoubleType, false) :: Nil)
  
  // for each loop on each class type
  classRecord.rdd.foreach{
    //filter students based on camps
    var campYellow =studentRecord.filter($"Camp" === "yellow")
    var campWhite =studentRecord.filter($"Camp" === "white")
    var campRed =studentRecord.filter($"Camp" === "red")
  
    // since I know that calculation for first class is average, so representing calculation only for class first
    val avgcampYellow  =  averageOnName(campYellow)
    val avgcampWhite   =  averageOnName(campWhite)
    val avgcampRed   =  averageOnName(campRed)
  
    // union of all
    val rddYellow = sc.parallelize (avgcampYellow).map (x =&gt; org.apache.spark.sql.Row(x._1, x._2.asInstanceOf[Number].doubleValue()))
    //conversion of rdd to frame
    var dfYellow = sqlContext.createDataFrame(rddYellow, schema)
    //union with yellow camp data
    val rddWhite = sc.parallelize (avgcampWhite).map (x =&gt; org.apache.spark.sql.Row(x._1, x._2.asInstanceOf[Number].doubleValue()))
    //conversion of rdd to frame
    var dfWhite = sqlContext.createDataFrame(rddWhite, schema)
    var dfYellWhite = dfYellow.union(dfWhite)
    //union with yellow,white camp data
    val rddRed = sc.parallelize (avgcampRed).map (x =&gt; org.apache.spark.sql.Row(x._1, x._2.asInstanceOf[Number].doubleValue()))
    //conversion of rdd to frame
    var dfRed = sqlContext.createDataFrame(rddRed, schema)
    var dfYellWhiteRed = dfYellWhite .union(dfRed)
    // other modifications and final result to hive
  }
  </code></pre>
  
  <p>Here I am struggling with: </p>
  
  <ol>
  <li>Hardcoding yellow, red and white, there may be additional camp types as well. </li>
  <li>The dataframe is currently being filtered many times which could be improved.</li>
  <li>I'm not able to figure out how to calculate differently according to class calculation type (i.e. use sum/averge depending on the class type).</li>
  </ol>
  
  <p>Any help is appreciated.</p>