title: Pyspark pyarrow pandas_udf - GROUPED_MAP return dataframe with None NaN for IntegerType, TimestampType
body: <p>Best</p>
  
  <p>At this moment I'm experimenting with pyspark <strong>pandas_udf</strong>, but unfortunately I encounter some issues when I return a DataFrame with: NA, None or NaNs in it. If I'm using a FloatType then the result is okey but as soon as I'm using a IntegerType, TimestampType etc... I receive an error &amp; it doesn't work any more.</p>
  
  <p>Here are some examples of what works &amp; what doesn't work:</p>
  
  <p><code>What does work?</code><br/>
  Example 1)</p>
  
  <pre><code>custom_schema = StructType([
                          StructField('User',StringType(),True),
                          StructField('Sport',StringType(),True),
                          StructField('Age',IntegerType(),True),
                          StructField('Age_lag',FloatType(),True),
                          ])
  
  # the schema is what it needs as an output format
  @pandas_udf(custom_schema, PandasUDFType.GROUPED_MAP)
  def my_custom_function(pdf):
      # Input/output are both a pandas.DataFrame
  
      #return a totalaly different DataFrame
      dt = pd.DataFrame({'User': ['Alice', 'Bob'], 'Sport': ['Football', 'Basketball'], 'Age': [27, 34]})
      dt['Age_lag'] = dt['Age'].shift(1)
  
      return dt
  
  df.groupby('id').apply(my_custom_function).toPandas() 
  </code></pre>
  
  <p>Result:</p>
  
  <pre><code>    User    Sport   Age     Age_lag
  0   Alice   Football    27  NaN
  1   Bob     Basketball  34  27.0
  2   Alice   Football    27  NaN
  3   Bob     Basketball  34  27.0
  </code></pre>
  
  <p>Example 2)</p>
  
  <p>If we change the <strong>Type</strong> of <em>Age_lag</em> into <em>IntegerType()</em> &amp; fill the Na's by -1 then we still have a valid result (no NaNs)</p>
  
  <pre><code>custom_schema = StructType([
                          StructField('User',StringType(),True),
                          StructField('Sport',StringType(),True),
                          StructField('Age',IntegerType(),True),
                          StructField('Age_lag',IntegerType(),True),
                          ])
  
  # the schema is what it needs as an output format
  @pandas_udf(custom_schema, PandasUDFType.GROUPED_MAP)
  def my_custom_function(pdf):
      # Input/output are both a pandas.DataFrame
  
      #return a totalaly different DataFrame
      dt = pd.DataFrame({'User': ['Alice', 'Bob'], 'Sport': ['Football', 'Basketball'], 'Age': [27, 34]})
      dt['Age_lag'] = dt['Age'].shift(1).fillna(-1)
  
      return dt
  
  df.groupby('id').apply(my_custom_function).toPandas() 
  </code></pre>
  
  <p>Result:</p>
  
  <pre><code>    User    Sport   Age     Age_lag
  0   Alice   Football    27  -1
  1   Bob     Basketball  34  27
  2   Alice   Football    27  -1
  3   Bob     Basketball  34  27
  </code></pre>
  
  <p><br/><br/></p>
  
  <p>What does not work?</p>
  
  <p>Example 3)</p>
  
  <p><strong>If</strong> we leave out the <strong>.fillna(-1)</strong> then we receive the next error</p>
  
  <pre><code>custom_schema = StructType([
                          StructField('User',StringType(),True),
                          StructField('Sport',StringType(),True),
                          StructField('Age',IntegerType(),True),
                          StructField('Age_lag',IntegerType(),True),
                          ])
  
  # the schema is what it needs as an output format
  @pandas_udf(custom_schema, PandasUDFType.GROUPED_MAP)
  def my_custom_function(pdf):
      # Input/output are both a pandas.DataFrame
  
      #return a totalaly different DataFrame
      dt = pd.DataFrame({'User': ['Alice', 'Bob'], 'Sport': ['Football', 'Basketball'], 'Age': [27, 34]})
      dt['Age_lag'] = dt['Age'].shift(1)
  
      return dt
  
  df.groupby('id').apply(my_custom_function).toPandas() 
  </code></pre>
  
  <p>Result: <strong>pyarrow.lib.ArrowInvalid: Floating point value truncated</strong></p>
  
  <p><br/>  <br/>
  Example 4)</p>
  
  <p>Last but not least, if we just send a static dataframe back where age_lag contains a <strong>None</strong>, then it doesn't work either.</p>
  
  <pre><code>from pyspark.sql.types import StructType,NullType, StructField,FloatType, LongType, DoubleType, StringType, IntegerType
  # true means,  accepts nulls
  custom_schema = StructType([
                          StructField('User',StringType(),True),
                          StructField('Sport',StringType(),True),
                          StructField('Age',IntegerType(),True),
                          StructField('Age_lag',IntegerType(),True),
                          ])
  
  # the schema is what it needs as an output format
  @pandas_udf(custom_schema, PandasUDFType.GROUPED_MAP)
  def my_custom_function(pdf):
      # Input/output are both a pandas.DataFrame
  
      #return a totalaly different DataFrame
      dt = pd.DataFrame({'User': ['Alice', 'Bob'], 
                        'Sport': ['Football', 'Basketball'], 
                          'Age': [27, 34], 
                      'Age_lag': [27, None]})
  
  
      return dt
  
  df.groupby('id').apply(my_custom_function).toPandas() 
  </code></pre>
  
  <p><strong>Questions</strong>: </p>
  
  <ul>
  <li>How do you deal with this?</li>
  <li>Is it a poor design ?
  
  <ul>
  <li>(because I can imagine 1000 cases where I do want to return NaNs and Nones)</li>
  </ul></li>
  <li>Do we really have to fill in all the missing values? and replace them back afterwards? or work with floats instead if integers?  etc? </li>
  <li>will this be solved in the near future? (because pandas_udf s are quite new)</li>
  </ul>