title: Adding a group count column to a PySpark dataframe
body: <p>I am coming from R and the <a href="https://www.tidyverse.org/" rel="noreferrer">tidyverse</a> to PySpark due to its superior Spark handling, and I am struggling to map certain concepts from one context to the other.</p>
  
  <p>In particular, suppose that I had a dataset like the following</p>
  
  <pre><code>x | y
  --+--
  a | 5
  a | 8
  a | 7
  b | 1
  </code></pre>
  
  <p>and I wanted to add a column containing the number of rows for each <code>x</code> value, like so:</p>
  
  <pre><code>x | y | n
  --+---+---
  a | 5 | 3
  a | 8 | 3
  a | 7 | 3
  b | 1 | 1
  </code></pre>
  
  <p>In dplyr, I would just say:</p>
  
  <pre class="lang-r prettyprint-override"><code>import(tidyverse)
  
  df &lt;- read_csv("...")
  df %&gt;%
      group_by(x) %&gt;%
      mutate(n = n()) %&gt;%
      ungroup()
  </code></pre>
  
  <p>and that would be that. I can do something almost as simple in PySpark if I'm looking to <strong>summarize</strong> by number of rows:</p>
  
  <pre class="lang-py prettyprint-override"><code>from pyspark.sql import SparkSession
  from pyspark.sql.functions import col
  
  spark = SparkSession.builder.getOrCreate()
  
  spark.read.csv("...") \
      .groupBy(col("x")) \
      .count() \
      .show()
  </code></pre>
  
  <p>And I thought I understood that <code>withColumn</code> was equivalent to dplyr's <code>mutate</code>. However, when I do the following, PySpark tells me that <code>withColumn</code> is not defined for <code>groupBy</code> data:</p>
  
  <pre class="lang-py prettyprint-override"><code>from pyspark.sql import SparkSession
  from pyspark.sql.functions import col, count
  
  spark = SparkSession.builder.getOrCreate()
  
  spark.read.csv("...") \
      .groupBy(col("x")) \
      .withColumn("n", count("x")) \
      .show()
  </code></pre>
  
  <p>In the short run, I can simply create a second dataframe containing the counts and join it to the original dataframe. However, it seems like this could become inefficient in the case of large tables. What is the canonical way to accomplish this?</p>