Bug: int columns produce inconsistent results on pandas UDFs
Status: Open | Priority: Major | Resolution: 
Created: 13/Aug/20 19:59
Affects Version/s: 3.0.0

This is similar to SPARK-30187 but I personally consider this data corruption.

If I have a simple pandas UDF
{code}
 >>> def add(a, b):
        return a + b
 >>> my_udf = pandas_udf(add, returnType=LongType())
{code}

And I want to process some data with it, say 32 bit ints
{code}
>>> df = spark.createDataFrame([(1037694399, 1204615848),(3,4)], StructType([StructField("a", IntegerType()), StructField("b", IntegerType())]))

>>> df.select(my_udf(col("a") - 3, col("b")).show()
+----------+----------+---------------+
|         a|         b|add((a - 3), b)|
+----------+----------+---------------+
|1037694399|1204615848|    -2052657052|
|         3|         4|              4|
+----------+----------+---------------+
{code}

I get an integer overflow for the data as I would expect.  But as soon as I add a {{None}} to the data, even on a different row the result I get back is totally different.

{code}
>>> df = spark.createDataFrame([(1037694399, 1204615848),(3,None)], StructType([StructField("a", IntegerType()), StructField("b", IntegerType())]))

>>> df.select(col("a"), col("b"), my_udf(col("a") - 3, col("b"))).show()
+----------+----------+---------------+
|         a|         b|add((a - 3), b)|
+----------+----------+---------------+
|1037694399|1204615848|     2242310244|
|         3|      null|           null|
+----------+----------+---------------+
{code}

The integer overflow disappears.  This is because arrow and/or pandas changes the data type to a float in order to be able to store the null value.  So then the processing is being done on floating point there is no overflow.  This in and of itself is annoying but understandable because it is dealing with a limitation in pandas. 

Where it becomes a bug is that this happens on a per batch basis.  This means that I can have the same two rows in different parts of my data set and get different results depending on their proximity to a null value.

{code}
>>> df = spark.createDataFrame([(1037694399, 1204615848),(3,None),(1037694399, 1204615848),(3,4)], StructType([StructField("a", IntegerType()), StructField("b", IntegerType())]))
>>> df.select(col("a"), col("b"), my_udf(col("a") - 3, col("b"))).show()
+----------+----------+---------------+
|         a|         b|add((a - 3), b)|
+----------+----------+---------------+
|1037694399|1204615848|     2242310244|
|         3|      null|           null|
|1037694399|1204615848|     2242310244|
|         3|         4|              4|
+----------+----------+---------------+

>>> spark.conf.set('spark.sql.execution.arrow.maxRecordsPerBatch', '2')
>>> df.select(col("a"), col("b"), my_udf(col("a") - 3, col("b"))).show()
+----------+----------+---------------+
|         a|         b|add((a - 3), b)|
+----------+----------+---------------+
|1037694399|1204615848|     2242310244|
|         3|      null|           null|
|1037694399|1204615848|    -2052657052|
|         3|         4|              4|
+----------+----------+---------------+

{code}

For me personally I would prefer to have all nullable integer columns upgraded to float all the time, that way it is at least consistent.