title: Pyspark: How to Apply UDF only on Rows with NotNull Values
body: <p>I have a pyspark dataframe and would like to apply an UDF on a column with Null values.</p>
  <p>Below is my dataframe:</p>
  <pre><code>+----+----+
  |   a|   b|
  +----+----+
  |null|  00|
  |.Abc|null|
  |/5ee|  11|
  |null|   0|
  +----+----+
  </code></pre>
  <p>Below is the desired dataframe (remove punctuations and change string values to upper case in column <strong>a</strong> if row values are not Null):</p>
  <pre><code>+----+----+
  |   a|   b|
  +----+----+
  |null|  00|
  | ABC|null|
  | 5EE|  11|
  |null|   0|
  +----+----+
  </code></pre>
  <p>Below is my UDF and code:</p>
  <pre><code>import pyspark.sql.functions as F
  import re
  
  remove_punct = F.udf(lambda x: re.sub('[^\w\s]', '', x))
  df = df.withColumn('a', F.when(F.col(&quot;a&quot;).isNotNull(), F.upper(remove_punct(F.col(&quot;a&quot;)))))
  </code></pre>
  <p>Below is the error:</p>
  <pre><code>TypeError: expected string or bytes-like object
  </code></pre>
  <p>Can you please suggest what would be the optimal solution the get the desired DF?</p>
  <p>Thanks in advance!</p>