title: Encode and assemble multiple features in PySpark
body: <p>I have a Python class that I'm using to load and process some data in Spark. Among various things I need to do, I'm generating a list of dummy variables derived from various columns in a Spark dataframe. My problem is that I'm not sure how to properly define a User Defined Function to accomplish what I need.</p>
  
  <p>I <em>do</em> currently have a method that, when mapped over the underlying dataframe RDD, solves half the problem (remember that this is a method in a larger <code>data_processor</code> class):</p>
  
  <pre><code>def build_feature_arr(self,table):
      # this dict has keys for all the columns for which I need dummy coding
      categories = {'gender':['1','2'], ..}
  
      # there are actually two differnt dataframes that I need to do this for, this just specifies which I'm looking at, and grabs the relevant features from a config file
      if table == 'users':
          iter_over = self.config.dyadic_features_to_include
      elif table == 'activty':
          iter_over = self.config.user_features_to_include
  
      def _build_feature_arr(row):
          result = []
          row = row.asDict()
          for col in iter_over:
              column_value = str(row[col]).lower()
              cats = categories[col]
              result += [1 if column_value and cat==column_value else 0 for cat in cats]
          return result
      return _build_feature_arr
  </code></pre>
  
  <p>Essentially what this does is, for the specified dataframe, takes the categorical variable values for the specified columns, and returns a list of the values of these new dummy variables. That means the following code:</p>
  
  <pre><code>data = data_processor(init_args)
  result = data.user_data.rdd.map(self.build_feature_arr('users'))
  </code></pre>
  
  <p>returns something like:</p>
  
  <pre><code>In [39]: result.take(10)
  Out[39]:
  [[1, 0, 0, 0, 1, 0],
   [1, 0, 0, 1, 0, 0],
   [1, 0, 0, 0, 0, 0],
   [1, 0, 1, 0, 0, 0],
   [1, 0, 0, 1, 0, 0],
   [1, 0, 0, 1, 0, 0],
   [0, 1, 1, 0, 0, 0],
   [1, 0, 1, 1, 0, 0],
   [1, 0, 0, 1, 0, 0],
   [1, 0, 0, 0, 0, 1]]
  </code></pre>
  
  <p>This is exactly what I want in terms of generating the list of dummy variables I want, but here's my question: How can I either (a) make a UDF with similar functionality that I can use in a Spark SQL query (or some other way, I suppose), or (b) take the RDD resulting from the map described above and add it as a new column to the user_data dataframe?</p>
  
  <p>Either way, what I need to do is generate a new dataframe containing the columns from user_data, along with a new column (let's call it <code>feature_array</code>) containing the output of the function above (or something functionally equivalent).</p>