title: Pandas on spark dataframe assign function ValueError
body: <p>I want to create a new column in a pandas-on-spark dataframe. In previous attempts, I tried pure pandas syntax like `df['column']=...' but I struggle against pickling error which I could not understand. So I find out a specific method of pandas dataframe that should do the trick. Unfortunately, I get some strange ValueError.</p>
  <pre><code>import pyspark.pandas as ps
  from pyspark.pandas import DataFrame
  from functools import partial
  from collections import defaultdict
  from typing import Dict, List
  
  
  def chars_to_main_df(dataset_df: DataFrame, char_score_map: Dict[str, float]):
  
      chars = dataset_df.char
      chars_scores = chars.apply(partial(chars_to_main, char_score_map=char_score_map))
  
      return chars_scores
  
  
  def chars_to_main(chars: List[str], char_score_map: Dict[str, float]):
  
      mapped = defaultdict(float)
      for cat in chars:
          mapped[cat] += char_score_map[cat]
  
      return dict(mapped)
  
  
  dict_ = {&quot;char&quot;: [[&quot;a&quot;], [&quot;b&quot;, &quot;c&quot;], [&quot;d&quot;]]}
  df = ps.DataFrame(data=dict_, columns=[&quot;char&quot;])
  char_score_map = {&quot;a&quot;: 1.0, &quot;b&quot;: 2.0, &quot;c&quot;: 3.0, &quot;d&quot;: 4.0}
  X = df.assign(cat_scores=partial(chars_to_main_df, char_score_map=char_score_map))
  X.to_pandas()
  </code></pre>
  <p>I am not sure the problem is <code>assign</code>, it could also a wrong results of <code>apply</code> but I find out really hard to use debug to see the content of <code>chars_scores</code></p>
  <p>Input (df):</p>
  <pre><code>---------
       char                                                                       
  0     [a]
  1  [b, c]
  2     [d]
  ---------
  </code></pre>
  <p>Output (X):</p>
  <pre><code>     char                                                                       
  0     [{a:1}]
  1  [{b:2, c:3}]
  2     [{d:4}]
  </code></pre>