title: Efficient way of applying specific function to structured column in Spark data frame?
body: <p>I have data in a Spark data frame, with a column <code>col</code> that contains <em>structured data</em> of the form:</p>
  
  <pre><code>------ col -------   # Column whose elements are structures
  field0  field1 …     # StructType with StructFields (variable names and count)
  [1,2,3] [4,5]  [6]   # Each field is of type ArrayType
  [1,2]   [3]    []
  …
  </code></pre>
  
  <p>where <em>the number and the names of the fields are not fixed</em>.</p>
  
  <p><strong>What is the most efficient way of calculating the total number of elements in each row?</strong> In the above example, the expected resulting data frame would thus be:</p>
  
  <pre><code>num_elements
  6
  3
  …
  </code></pre>
  
  <p>There is always the solution of a user defined function:</p>
  
  <pre><code>from pyspark.sql.types import IntegerType
  
  def num_elements(all_arrays_in_row):
      return sum(map(len, all_arrays_in_row))
  num_elements = pyspark.sql.functions.udf(num_elements, IntegerType())
  
  data_frame.select(num_elements(data_frame.col)).show()  # Number of elements in each row
  </code></pre>
  
  <p>Now, I am not sure whether this is generally efficient, because:</p>
  
  <ol>
  <li>Function <code>num_elements()</code> is in Python.</li>
  <li>If the fields happen to not be stored together for some reason, the <code>map()</code> forces a fetch of each array before calculating their length.</li>
  </ol>
  
  <p>More generally, a "pure" Spark approach would be more efficient, but it is eluding me. What I tried so far is the following, but this is way more cumbersome than the approach above, and is also not complete:</p>
  
  <ol>
  <li>Get the field names <code>field0</code>, etc. with <code>[field.name for field in data_frame.select("col").schema.fields[0].dataType.fields]</code> (cumbersome).</li>
  <li><p>For each field name, efficiently calculate the size of its array:</p>
  
  <pre><code>sizes_one_field = data_frame.select(pyspark.sql.functions.size(
                                      data_frame.col.getField(field_name))
  </code></pre></li>
  </ol>
  
  <p>Now, I am stuck at this point because I am not sure how to sum together the 1-column data frames <code>sizes_one_field</code> (there is one for each field name). Plus, maybe there is a way of directly applying the <code>size()</code> function to each field of column <code>col</code> in Spark (through some kind of map?)? Or some completely different approach to getting the total number of elements in each row?</p>