Pyspark / DataBricks DataFrame size estimation

pyspark_tricks.py

from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
def _to_java_object_rdd(rdd):  
    """ Return a JavaRDD of Object by unpickling
    It will convert each Python object into Java object by Pyrolite, whenever the
    RDD is serialized in batc h or not.
    """
    rdd = rdd._reserialize(AutoBatchedSerializer(PickleSerializer()))
    return rdd.ctx._jvm.org.apache.spark.mllib.api.python.SerDe.pythonToJava(rdd._jrdd, True)

def estimate_df_size(df):
    JavaObj = _to_java_object_rdd(df.rdd)
    nbytes = sc._jvm.org.apache.spark.util.SizeEstimator.estimate(JavaObj)
    return nbytes