code: def saveAsTFRecords(df, output_dir):
    
    tf_rdd = df.rdd.mapPartitions(toTFExample(df.dtypes))
    tf_rdd.saveAsNewAPIHadoopFile(output_dir, "org.tensorflow.hadoop.io.TFRecordFileOutputFormat",
                                  keyClass="org.apache.hadoop.io.BytesWritable",
                                  valueClass="org.apache.hadoop.io.NullWritable")
docs: Save a Spark DataFrame as TFRecords.
  
    This will convert the DataFrame rows to TFRecords prior to saving.
  
    Args:
      :df: Spark DataFrame
      :output_dir: Path to save TFRecords