title: Group by key error for dask dataframe map_partitions
body: <p>I'm trying to run a dask job on a big dataframe. Essentially there are two groups within the dataframe, and I'm performing a fuzzy match between them. One group is the group doing the matching, and the other group is the one being matched. The matching group has value 3.0, and the matched group has value in either 1.5 or 2.0. </p>
  
  <p>My code works for the most part, but after its completion (which takes a couple hours), it gives me a key error for the matching group. </p>
  
  <p>Here is my code:</p>
  
  <pre><code>for i in [2.0,1.5]:
      FuzzyWuzzyResults = emb.map_partitions(
          lambda df: df.groupby('phase').get_group(3.0)['drugs'].apply(
              fuzzy_match, args=(df.groupby('phase').get_group(i)['drugs'],
              fuzz.token_sort_ratio,85)), meta=('results')
          ).compute() 
  </code></pre>
  
  <p>It'll work for most partitions of the huge pandas dataframe, but at the end, I get a couple "KeyError 3.0" errors, and the code breaks in between. How can I fix this? If possible, I would not want to use groupby and rather break the og dataframe into two and apply the dask job that way, but I'm not sure.</p>