From Pandas groupBy to PySpark groupBy - pandas

Question

Consider an Spark DataFrame, wherein we have few columns. The goal is to perform groupBy operation on it without converting it to Pandas DataFrame. An equivalent Pandas groupBy code looks something like this:
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()
And I'm intending to write this in PySpark. So far I have come up with something like this in PySpark:
gdf = df.groupBy([
'col_1',
'col_2'
]).agg({
'c': 'avg',
'b': 'sum'
}).withColumnRenamed('sum(b)', 'new_b')
However, I am not sure how to go about 'a': x['a'].values[0] and 'cnt': len(x). I thought about using collect_list from from pyspark.sql import functions but that slaps my face with Column object is not Callable. Any idea how to accomplish the aforementioned conversion? Thanks!
[UPDATE] Would it make sense to perform count operation on any column in order to get cnt? Say I do this:
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')

Answer 1

I have this toy solution using PySpark function sum, avg, count and first. note that I use Spark 2.1 in this solution. Hope this help a bit!
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