pyspark批量生成tfrecord文件

人狠话不多，直接上代码！

from pyspark import SparkConf
from pyspark.sql import SparkSession
from pyspark.sql import HiveContext
from pyspark.sql import Row
from pyspark import SparkFiles
from pyspark.sql.types import *

sqlContext = HiveContext(spark)
sqlContext.sql("set hive.exec.dynamic.partition=true")
sqlContext.sql("set hive.exec.dynamic.partition.mode=strick")

def parse_example(r):
    result = []  # 返回result，和schema对应list
    # TODO something
    result.append(feature)
    result.append(label)
    return result

# 读hive数据
train_data = sqlContext.sql("SELECT feature,label from db.my_table")

# 定义数据结构
fields = [
    StructField('feature', ArrayType(IntegerType(), True)),
    StructField("label", ArrayType(IntegerType(), True))
]
schema = StructType(fields)

# 样本处理
result_rdd = train_data.rdd.repartition(3000).map(parse_example)
df = spark.createDataFrame(result_rdd, schema)
tfrecord_output_file = '你的文件存储路径'
file_num = 100  #定义生成多少个文件
df.repartition(file_num) \
        .write \
        .format("tfrecords") \
        .option("recordType", "Example") \
        .option("codec", "org.apache.hadoop.io.compress.DefaultCodec") \
        .save(tfrecord_output_file, mode="overwrite")