title: Spark dataframe to numpy array via udf or without collecting to driver
body: <p>Real life df is a massive dataframe that cannot be loaded into driver memory.
  Can this be done using regular or pandas udf?</p>
  
  <pre><code># Code to generate a sample dataframe
  
  from pyspark.sql import functions as F
  from pyspark.sql.types import *
  import pandas as pd
  import numpy as np
  
  sample = [['123',[[0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1], [0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
        ['345',[[1,0,0,0,0,1,1,1,0,1,1,0,1,0,0,0,1,1,1,1,1,1], [0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
        ['425',[[1,1,0,0,0,1,0,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1],[0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
        ]
  
  df = spark.createDataFrame(sample,["id", "data"])
  </code></pre>
  
  <p>Here's the logic that needs to be parallelized without relying on driver memory.</p>
  
  <p>Input: Spark dataframe
  Output: numpy array to be fed into horovod (Something like this: <a href="https://docs.databricks.com/applications/deep-learning/distributed-training/mnist-tensorflow-keras.html" rel="nofollow noreferrer">https://docs.databricks.com/applications/deep-learning/distributed-training/mnist-tensorflow-keras.html</a>)</p>
  
  <pre><code>pandas_df = df.toPandas() # Not possible in real life
  data_array = np.asarray(list(pandas_df.data.values))
  data_array = data_array.reshape(data_array.shape[0], data_array.shape[1], -1, 1, order='F')
  data_array = data_array.reshape(data_array.shape[0],data_array.shape[1],-1,1,1,order="F").transpose(0,1,3,2,-1)
  # Some more numpy specific transformations ..
  </code></pre>
  
  <p>Here's an approach that didn't work:</p>
  
  <pre><code>@pandas_udf(ArrayType(IntegerType()), PandasUDFType.SCALAR)
  def generate_feature(x):
      data_array = np.asarray(x)
      data_array = data_array.reshape(data_array.shape[0], ..
      ...
      return pd.Series(data_array)
  
  df = df.withColumn("data_array", generate_feature(df.data))
  </code></pre>