人狠话不多,直接上代码!
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")