Applying UDFs on GroupedData in PySpark (with functioning python example)
<p>I have this python code that runs locally in a pandas dataframe:</p>

<pre><code>df_result = pd.DataFrame(df
                          .groupby('A')
                          .apply(lambda x: myFunction(zip(x.B, x.C), x.name))
</code></pre>

<p>I would like to run this in PySpark, but having trouble dealing with pyspark.sql.group.GroupedData object.</p>

<p>I've tried the following:</p>

<pre><code>sparkDF
 .groupby('A')
 .agg(myFunction(zip('B', 'C'), 'A')) 
</code></pre>

<p>which returns</p>

<pre><code>KeyError: 'A'
</code></pre>

<p>I presume because 'A' is no longer a column and I can't find the equivalent for x.name.</p>

<p>And then</p>

<pre><code>sparkDF
 .groupby('A')
 .map(lambda row: Row(myFunction(zip('B', 'C'), 'A'))) 
 .toDF()
</code></pre>

<p>but get the following error:</p>

<pre><code>AttributeError: 'GroupedData' object has no attribute 'map'
</code></pre>

<p>Any suggestions would be really appreciated!</p>

<p>What you are trying to is write a UDAF (User Defined Aggregate Function) as opposed to a UDF (User Defined Function).  UDAFs are functions that work on data grouped by a key.  Specifically they need to define how to merge multiple values in the group in a single partition, and then how to merge the results across partitions for key.  There is currently no way in python to implement a UDAF, they can only be implemented in scala.  </p>

<p>But, you can work around it in Python.  You can use collect set to gather your grouped values and then use a regular UDF to do what you want with them.  The only caveat is collect_set only works on primative values, so you will need to encode them down to a string.</p>

<pre><code>from pyspark.sql.types import StringType
from pyspark.sql.functions import col, collect_set, concat_ws, udf

def myFunc(data_list):
    for val in data_list:
        b, c = data.split(',')
        # do something

    return &lt;whatever&gt;

myUdf = udf(myFunc, StringType())

df.withColumn('data', concat_ws(',', col('B'), col('C'))) \
  .groupBy('A').agg(collect_list('data').alias('data'))
  .withColumn('data', myUdf('data'))
</code></pre>

<p>Use collect_set if you want deduping.  Also, if you have lots of values for some of your keys, this will be slow because all values for a key will need to be collected in a single partition somewhere on your cluster.  If your end result is a value you build by combining the values per key in some way (for example summing them) it might be faster to implement it using the <a href="http://spark.apache.org/docs/1.6.2/api/python/pyspark.html#pyspark.RDD.aggregateByKey" rel="nofollow">RDD aggregateByKey</a> method which lets you build an intermediate value for each key in a partition before shuffling data around.</p>

python apache-spark pyspark spark-dataframe