mapPartition function to convert a Spark RDD of Row into an RDD of serialized tf.train.Example bytestring.

  Note that tf.train.Example is a fairly flat structure with limited datatypes, e.g. tf.train.FloatList,
  tf.train.Int64List, and tf.train.BytesList, so most DataFrame types will be coerced into one of these types.

  Args:
    :dtypes: the DataFrame.dtypes of the source DataFrame.

  Returns:
    A mapPartition function which converts the source DataFrame into tf.train.Example bytestrings.

def toTFExample(dtypes):
  """mapPartition function to convert a Spark RDD of Row into an RDD of serialized tf.train.Example bytestring.

  Note that tf.train.Example is a fairly flat structure with limited datatypes, e.g. tf.train.FloatList,
  tf.train.Int64List, and tf.train.BytesList, so most DataFrame types will be coerced into one of these types.

  Args:
    :dtypes: the DataFrame.dtypes of the source DataFrame.

  Returns:
    A mapPartition function which converts the source DataFrame into tf.train.Example bytestrings.
  """
  def _toTFExample(iter):

    # supported type mappings between DataFrame.dtypes and tf.train.Feature types
    float_dtypes = ['float', 'double']
    int64_dtypes = ['boolean', 'tinyint', 'smallint', 'int', 'bigint', 'long']
    bytes_dtypes = ['binary', 'string']
    float_list_dtypes = ['array<float>', 'array<double>']
    int64_list_dtypes = ['array<boolean>', 'array<tinyint>', 'array<smallint>', 'array<int>', 'array<bigint>', 'array<long>']

    def _toTFFeature(name, dtype, row):
      feature = None
      if dtype in float_dtypes:
        feature = (name, tf.train.Feature(float_list=tf.train.FloatList(value=[row[name]])))
      elif dtype in int64_dtypes:
        feature = (name, tf.train.Feature(int64_list=tf.train.Int64List(value=[row[name]])))
      elif dtype in bytes_dtypes:
        if dtype == 'binary':
          feature = (name, tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(row[name])])))
        else:
          feature = (name, tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(row[name]).encode('utf-8')])))
      elif dtype in float_list_dtypes:
        feature = (name, tf.train.Feature(float_list=tf.train.FloatList(value=row[name])))
      elif dtype in int64_list_dtypes:
        feature = (name, tf.train.Feature(int64_list=tf.train.Int64List(value=row[name])))
      else:
        raise Exception("Unsupported dtype: {0}".format(dtype))
      return feature

    results = []
    for row in iter:
      features = dict([_toTFFeature(name, dtype, row) for name, dtype in dtypes])
      example = tf.train.Example(features=tf.train.Features(feature=features))
      results.append((bytearray(example.SerializeToString()), None))
    return results

  return _toTFExample