title: pandas_udf failing for large dataset
body: <p>I have a problem with a large dataframe (40 billions+ rows) where I am creating a new column passing an array column to a UDF. A PySpark UDF works for a smaller dataset and not working for more than few thousand records. So I am trying pandas_udf with apache arrow.</p>
  <pre><code>group_key col1 col2
  123       a    5
  123       a    6
  123       b    6
  123       cd   3 
  123       d    2
  123       ab   9
  456       d    4  
  456       ad   6 
  456       ce   7 
  456       a    4 
  456       s    3 
  </code></pre>
  <p>Desired output</p>
  <pre><code>group_key output
  123       9.2
  456       7.3
  </code></pre>
  <p>This is a sample code, where I am passing an array of arrays to the UDF. arr_of_arr in func is showing as numpy ndarray. I return StringType() here.</p>
  <pre><code>df_arr2=df_new2.groupBy(&quot;group_key&quot;).agg(F.collect_list( F.array(&quot;col1&quot;, col2&quot;)).alias(&quot;key_value_arr&quot;))
      
  
  |group_key|key_value_arr                                                                                                                              
  |123       |[[a, 5], [a, 6], [b, 6], [cd, 3], [d 2], [ab, 9]]|
  |456       |[[d, 4], [ad, 6], [ce, 7], [a, 4], [s, 3]]|
  
  
  
  df_arr2.printSchema()
  
  @pandas_udf(StringType())
  def func(arr_of_arr):
      res=''
      # arr_of_arr.values is 'numpy.ndarray' 
      
      for index, value in np.ndenumerate(arr_of_arr):  ## arr_of_arr has only 1 row with array objects . 0 is the index , value is the array of arrays objects
          for index1, value1 in value.tolist():
              res=str(value1) ## Just returning the last value here
      return pd.Series(res)
  
  df_arr2.withColumn('col3',func(df_arr2.key_value_arr)).show(truncate=False)
  </code></pre>
  <p>This code returns correct value from the inside array elements like &quot;5&quot;, &quot;7&quot; etc. for few records. For full dataset it errors probably performance related .</p>
  <blockquote>
  <pre><code>An error occurred while calling o480.showString.
    File &quot;/hadoop/3/yarn/local/usercache/b_incdata_rw/appcache/application_1660704390900_2339796/container_e3797_1660704390900_2339796_01_000001/py4j-0.10.7-src.zip/py4j/protocol.py&quot;,    
  </code></pre>
  <p>line 328, in get_return_value
  format(target_id, &quot;.&quot;, name), value)
  Py4JJavaError: An error occurred while calling o480.showString.
  : org.apache.spark.SparkException: Job aborted due to stage failure: ResultStage  128</p>
  </blockquote>
  <pre><code>sample output:
  root
   |-- group_key: long (nullable = true)
   |   |-- key_value_arr: array (nullable = true)
   |    |-- element: array (containsNull = true)
   |    |    |-- element: string (containsNull = true)
  
  group_key|key_value_arr|col3|
  |123       |[[a, 5], [a, 6], [b, 6], [cd, 3], [d, 2], [ab, 9]]|9|
  </code></pre>
  <p>The UDF logic is much more complex than this. But I need the values from each array element as key and value and loop through them. Performance does not look okay for such logic. Any suggestion ?</p>