title: UDF and python function in Pyspark
body: <p>I have a dataframe in Pyspark as:</p>
  
  <pre><code>listA = [(1,'AAA','USA'),(2,'XXX','CHN'),(3,'KKK','USA'),(4,'PPP','USA'),(5,'EEE','USA'),(5,'HHH','THA')]
  df = spark.createDataFrame(listA, ['id', 'name','country'])
  </code></pre>
  
  <p>and I have created a dictionary as:</p>
  
  <pre><code>thedict={"USA":"WASHINGTON","CHN":"BEIJING","DEFAULT":"KEY NOT FOUND"}
  </code></pre>
  
  <p>and Then I created a UDF to get the matching key values from dictionary.</p>
  
  <pre><code>def my_func(letter):
      if(thedict.get(letter) !=None):
          return thedict.get(letter)
      else:
          return thedict.get("DEFAULT")
  </code></pre>
  
  <p>I am getting below error when trying to call function as:</p>
  
  <pre><code>df.withColumn('CAPITAL',my_func(df.country))
  
    File "&lt;stdin&gt;", line 1, in &lt;module&gt;
    File "/usr/hdp/current/spark2-client/python/pyspark/sql/dataframe.py", line 1848, in withColumn
      assert isinstance(col, Column), "col should be Column"
  AssertionError: col should be Column
  </code></pre>
  
  <p>Whereas if I embedded it with pyspark.sql.functions, it's working fine.</p>
  
  <pre><code>from pyspark.sql.functions import col, udf
  udfdict = udf(my_func,StringType())
  
  df.withColumn('CAPITAL',udfdict(df.country)).show()
  
  +---+----+-------+-------------+
  | id|name|country|      CAPITAL|
  +---+----+-------+-------------+
  |  1| AAA|    USA|   WASHINGTON|
  |  2| XXX|    CHN|      BEIJING|
  |  3| KKK|    USA|   WASHINGTON|
  |  4| PPP|    USA|   WASHINGTON|
  |  5| EEE|    USA|   WASHINGTON|
  |  5| HHH|    THA|KEY NOT FOUND|
  +---+----+-------+-------------+
  </code></pre>
  
  <p>I couldn't understand what is the difference in these two calls? </p>