title: Weighted average with Spark Datasets without UDF
body: <p>While someone has already asked about computing a <a href="https://stackoverflow.com/questions/34325152/weighted-average-in-spark">Weighted Average in Spark</a>, in this question, I'm asking about using Datasets/DataFrames instead of RDDs.</p>
  
  <p>How do I compute a weighted average in Spark?  I have two columns: counts and previous averages:</p>
  
  <pre><code>case class Stat(name:String, count: Int, average: Double)
  val statset = spark.createDataset(Seq(Stat("NY", 1,5.0),
                             Stat("NY",2,1.5),
                             Stat("LA",12,1.0),
                             Stat("LA",15,3.0)))
  </code></pre>
  
  <p>I would like to be able to compute a weighted average like this:</p>
  
  <pre><code>display(statset.groupBy($"name").agg(sum($"count").as("count"),
                      weightedAverage($"count",$"average").as("average")))
  </code></pre>
  
  <p>One can use a UDF to get close:</p>
  
  <pre><code>val weightedAverage = udf(
    (row:Row)=&gt;{
      val counts = row.getAs[WrappedArray[Int]](0)
      val averages = row.getAs[WrappedArray[Double]](1)
      val (count,total) = (counts zip averages).foldLeft((0,0.0)){
        case((cumcount:Int,cumtotal:Double),(newcount:Int,newaverage:Double))=&gt;(cumcount+newcount,cumtotal+newcount*newaverage)}
      (total/count)  // Tested by returning count here and then extracting. Got same result as sum.
    }
  )
  
  display(statset.groupBy($"name").agg(sum($"count").as("count"),
                      weightedAverage(struct(collect_list($"count"),
                                      collect_list($"average"))).as("average")))
  </code></pre>
  
  <p>(Thanks to answers to <a href="https://stackoverflow.com/questions/41551410/passing-a-list-of-tuples-as-a-parameter-to-a-spark-udf-in-scala">Passing a list of tuples as a parameter to a spark udf in scala</a> for help in writing this)</p>
  
  <p>Newbies: Use these imports:</p>
  
  <pre><code>import org.apache.spark.sql._
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.types._
  import scala.collection.mutable.WrappedArray
  </code></pre>
  
  <p>Is there a way of accomplishing this with built-in column functions instead of UDFs? The UDF feels clunky and if the numbers get large you have to convert the Int's to Long's.</p>