title: Spark Dataframe grouping and partition by key with a set number of partitions.
body: <p>I have a spark dataframe with multiple labels and features coreesponding to each, like this:</p>
  
  <pre><code>+----------------+--------------------+
  |           label|       feature_paths|
  +----------------+--------------------+
  |         person1|[-0.015756417, 0....|
  |         person1|[-0.05177306, 0.1...|
  |         person1|[-0.11631858, 0.1...|
  |         person2|[-0.058303248, 0....|
  |         person2|[-0.03415013, 0.0...|
  +----------------+--------------------+
  </code></pre>
  
  <p>I want to train a clustering model for each label (person), so basically, I want to create an rdd for each label, and then run a map operation like <code>rdd.map(service)</code> which will eventually save a <code>gmm</code> model for each entity.  </p>
  
  <p>The code is like:</p>
  
  <pre><code>def service(rddentry):
  
      label = rddentry[0]
      features = rddentry[1]
  
      print(label)
  
      from sklearn.mixture import BayesianGaussianMixture
      from sklearn.externals import joblib
  
      gmm = BayesianGaussianMixture(n_components=3, covariance_type="diag", init_params='kmeans')
      model = gmm.fit(features)
      joblib.dump(model, str(label)+'.joblib') 
  
      return model
  </code></pre>
  
  <p>My goals, that I want to achieve is: </p>
  
  <ol>
  <li><p>Create an rdd where the number of partitions is equal to the number of unique labels, such that: rdd.getNumPartition() = no_of_unique_labels.
  Each rdd entry will have multiple features, belonging to a single label. </p></li>
  <li><p>Send each rdd partition to the service function. </p></li>
  </ol>
  
  <p>My experiments until now: </p>
  
  <ol>
  <li><p>When doing <code>sdf.repartition('label')</code>, it creates several empty dataframes. </p></li>
  <li><p><code>sdf.partionBy('label')</code> also does not work. It creates a random number of partitions. </p></li>
  </ol>
  
  <p>I have spent almost two days but of no concrete results until now. Any help or guidance in the right direction would be helpful. </p>