DataFrame simple aggregation performance benchmark

df.py

data = sqlContext.load("/home/rxin/ints.parquet")
data.groupBy("a").agg(col("a"), avg("num")).collect()

df.scala

val data = sqlContext.load("/home/rxin/ints.parquet")
data.groupBy("a").agg(col("a"), avg("num")).collect()

generatedata.py

import random
from pyspark.sql import Row

data = sc.parallelize(xrange(1000)).flatMap(lambda x: [Row(a=random.randint(1, 10), num=random.randint(1, 100), str=("a" * random.randint(1, 30))) for i in xrange(10000)])
dataTable = sqlContext.createDataFrame(data)
dataTable.saveAsParquetFile("/home/rxin/ints.parquet")

rdd.py

pdata = sqlContext.load("/home/rxin/ints.parquet").select("a", "num")

sum_count = (
  pdata.map(lambda x: (x.a, [x.num, 1]))
    .reduceByKey(lambda x, y:
		  [x[0] + y[0], x[1] + y[1]])
    .collect())
[(x[0], float(x[1][0]) / x[1][1]) for x in sum_count]

rdd.scala

val pdata = sqlContext.load("/home/rxin/ints.parquet").select("a", "num")

val sum_count = pdata.map { row => (row.getInt(0), (row.getInt(1), 1)) }
  .reduceByKey { (a, b) =>
    (a._1 + b._1, a._2 + b._2)
  }.collect()

sum_count.foreach { case (a, (sum, count)) => println(s"$a: ${sum/count}") }