Simulating UDAF on Pyspark for encapsulation
<p>I'm learning Spark with PySpark, and just hit a wall when trying to make things cleaner.</p>

<p>Say a have a dataframe that looks like this. (of course, with way more columns and rows)</p>

<pre><code>A | B |   C
--+---+------
a | 1 | 1.300
a | 2 | 2.500
a | 3 | 1.000
b | 1 | 120.0
b | 4 | 34.20
c | 2 | 3.442
</code></pre>

<p>and I want to run a bunch of <code>groupby -&gt; agg</code> on it, using basic <code>pyspark.sql.functions</code> , like <code>count()</code> and <code>mean()</code>, like this:</p>

<pre><code>df.groupby("A")\
    .agg(mean("B").alias("B_mean"),
         sum("C").alias("C_sum"),
         (countDistinct("B")/avg("C")).alias("New_metric"))
</code></pre>

<p>It works fine, runs relatively fast, and gives me the desired results. </p>

<p>But, eventually, slightly more complex functions will be needed, and, also, we want to make these easier to test. </p>

<p>How can one encapsulate these functions? Using <code>lambda</code>? Some way around UDFs?</p>

<p>I'm aware of UDAFs and that it's possible to write them in SCALA and import the code to PySpark, but, since all of our code base is already in Python, I would like to explore other options.</p>

<p>P.S.: We are running Spark 1.6.0</p>

<p>Function can be defined as a combination of <code>pyspark.sql.functions</code>:</p>

<ul>
<li><p>YES - go this way. For example:</p>

<pre><code>def sum_of_squares(col):
    return sum(col * col)

df.select(sum_of_squares(df["foo"]])

df.groupBy("foo").agg(sum_of_squares(df["bar"]])
</code></pre></li>
<li><p>NO - use RDD.</p></li>
</ul>

python apache-spark pyspark apache-spark-sql spark-dataframe