title: Counting words after grouping records (Part 2)
body: <p>Although I am <a href="https://stackoverflow.com/a/49925959/826983">having an answer</a> for what I want to achieve, the problem is that it's way to slow. The data set is not very large. It's 50GB in total but the affected part is probably just between 5 to 10GB of data. However, the following is what I require, but it's way to slow And by slow I mean it was running for an hour and it didn't terminate.</p>
  
  <pre class="lang-py prettyprint-override"><code>df_ = spark.createDataFrame([
      ('1', 'hello how are are you today'),
      ('1', 'hello how are you'),
      ('2', 'hello are you here'),
      ('2', 'how is it'),
      ('3', 'hello how are you'),
      ('3', 'hello how are you'),
      ('4', 'hello how is it you today')
  ], schema=['label', 'text'])
  
  tokenizer = Tokenizer(inputCol='text', outputCol='tokens')
  tokens = tokenizer.transform(df_)
  
  token_counts.groupby('label')\
      .agg(F.collect_list(F.struct(F.col('token'), F.col('count'))).alias('text'))\
      .show(truncate=False)
  </code></pre>
  
  <p>Which gives me the token count for each label:</p>
  
  <pre><code>+-----+----------------------------------------------------------------+
  |label|text                                                            |
  +-----+----------------------------------------------------------------+
  |3    |[[are,2], [how,2], [hello,2], [you,2]]                          |
  |1    |[[today,1], [how,2], [are,3], [you,2], [hello,2]]               |
  |4    |[[hello,1], [how,1], [is,1], [today,1], [you,1], [it,1]]        |
  |2    |[[hello,1], [are,1], [you,1], [here,1], [is,1], [how,1], [it,1]]|
  +-----+----------------------------------------------------------------+
  </code></pre>
  
  <p>However, I think the call to <code>explode()</code> is way too expensive for this.</p>
  
  <p>I don't know but it might be faster to count the tokens in each "dokument" and later merge it in a <code>groupBy()</code>:</p>
  
  <pre class="lang-py prettyprint-override"><code>df_.select(['label'] + [udf_get_tokens(F.col('text')).alias('text')])\
      .rdd.map(lambda x: (x[0], list(Counter(x[1]).items()))) \
      .toDF(schema=['label', 'text'])\
      .show() 
  </code></pre>
  
  <p>Gives the counts:</p>
  
  <pre><code>+-----+--------------------+
  |label|                text|
  +-----+--------------------+
  |    1|[[are,2], [hello,...|
  |    1|[[are,1], [hello,...|
  |    2|[[are,1], [hello,...|
  |    2|[[how,1], [it,1],...|
  |    3|[[are,1], [hello,...|
  |    3|[[are,1], [hello,...|
  |    4|[[you,1], [today,...|
  +-----+--------------------+
  </code></pre>
  
  <p>Is there a way to merge those token counts in a more efficient way?</p>