import functools


def balance_samples(spark, feature_df, label_col):
    """
    Create balanced sample of input data

    For a binary target column, sample the rows with replacement so as to create a
    synthetic data set with even representation for each class.  This expects the
    labeled column to have values of 0 or 1.

    Args:
        spark (``pyspark.sql.SparkSession``)
        feature_df (``pyspark.sql.DataFrame``): The source data
        label_col (str): The name of the column to rebalance the data by
    """
    counts = dict(map(lambda r: (r[label_col], r['count']), feature_df.groupBy(label_col).count().collect()))
    n_samples = functools.reduce(lambda a, b: a + b, counts.values())
    fractions = {1: float(n_samples)/counts[1] * 0.5, 0: float(n_samples)/counts[0] * 0.5}
    rdd = feature_df.rdd.map(lambda r: (r[label_col], r)).sampleByKey(True, fractions=fractions)
    return spark.createDataFrame(rdd.values(), feature_df.schema)