从Pandas group到PySpark groupBy

时间:2017-03-14 02:08:24

标签: pandas apache-spark pyspark

考虑Spark DataFrame,其中我们有几列。目标是在其上执行groupBy操作,而不将其转换为Pandas DataFrame。等效的Pandas groupBy代码如下所示:

def compute_metrics(x):
    return pd.Series({
        'a': x['a'].values[0],
        'new_b': np.sum(x['b']),
        'c': np.mean(x['c']),
        'cnt': len(x)
    })

data.groupby([
    'col_1',
    'col_2'
]).apply(compute_metrics).reset_index()

我打算在PySpark中写这个。到目前为止,我在PySpark

中提出了类似的内容
gdf = df.groupBy([
    'col_1',
    'col_2'
]).agg({
    'c': 'avg',
    'b': 'sum'
}).withColumnRenamed('sum(b)', 'new_b')

但是,我不确定如何处理'a': x['a'].values[0]'cnt': len(x)。我考虑过使用collect_list中的from pyspark.sql import functions,但是用Column object is not Callable打了我的脸。知道如何完成上述转换吗?谢谢!

[更新] 任何列上执行count操作是否有意义才能获得cnt?说我这样做:

gdf = df.groupBy([
    'col_1',
    'col_2'
]).agg({
    'c': 'avg',
    'b': 'sum',
    'some_column': 'count'
}).withColumnRenamed('sum(b)', 'new_b')
  .withColumnRenamed('count(some_column)', 'cnt')

1 个答案:

答案 0 :(得分:6)

我有这个玩具解决方案使用PySpark函数sumavgcountfirst注意我在此解决方案中使用Spark 2.1。希望这有点帮助!

from pyspark.sql.functions import sum, avg, count, first

# create toy example dataframe with column 'A', 'B' and 'C'
ls = [['a', 'b',3], ['a', 'b', 4], ['a', 'c', 3], ['b', 'b', 5]]
df = spark.createDataFrame(ls, schema=['A', 'B', 'C'])

# group by column 'A' and 'B' then performing some function here
group_df = df.groupby(['A', 'B'])
df_grouped = group_df.agg(sum("C").alias("sumC"), 
                          avg("C").alias("avgC"), 
                          count("C").alias("countC"), 
                          first("C").alias("firstC"))
df_grouped.show() # print out the spark dataframe