title: Pyspark udf doesn't work while Python function works
body: <p>I have a Python function:</p>
  <pre><code>def get_log_probability(string, transition_log_probabilities):
      string = ngrams(string, 2)
      terms = [transition_log_probabilities[bigram]
                         for bigram in string]
      log_probability = sum(terms)/len(terms) if len(terms) &gt; 0 else sum(terms)
      return log_probability
  </code></pre>
  <p>I want to use this function for Pyspark DataFrame column with <code>transition_log_probabilities</code> as a constant as follows:</p>
  <pre><code>transition_log_probabilities = {('a', 'a'): -3.688879454113936,
  ('a', 'b'): -3.688879454113936,
  ('a', 'c'): -3.688879454113936,
  ('b', 'a'): -3.688879454113936,
  ('b', 'b'): -3.688879454113936,
  ('b', 'c'): -3.688879454113936,
  ('c', 'a'): -3.688879454113936,
  ('c', 'b'): -3.688879454113936,
  ('c', 'c'): -3.688879454113936}
  </code></pre>
  <p>So I change it to Pyspark UDF:</p>
  <pre><code>def get_log_prob_udf(dictionary):
      return udf(lambda string: get_log_probability(string, dictionary), FloatType())
  </code></pre>
  <p>Even though  <code>get_log_probability(&quot;abc&quot;, transition_log_probabilities)</code> works and gives a result of <code>-3.688879454113936</code>, when I apply its UDF into Pyspark as follows:</p>
  <pre><code>df = df \
  .withColumn(&quot;string_log_probability&quot;, get_log_prob_udf(transition_log_probabilities)(col('string')))
  </code></pre>
  <p>It doesn't work and throws out the error of</p>
  <pre><code>An error occurred while calling o3463.showString.
  : org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 
  182.0 failed 1 times, most recent failure: Lost task 0.0 in stage 182.0 (TID 774) 
  (kubernetes.docker.internal executor driver): net.razorvine.pickle.PickleException: 
  expected zero arguments for construction of ClassDict (for numpy.dtype)
  </code></pre>
  <p>Does anyone know how to solve it? Thank you very much.</p>