[英]How to get rows from DF that contain value None in pyspark (spark)
[英]How do you get batches of rows from Spark using pyspark
我有一个包含超过 60 亿行数据的 Spark RDD,我想使用 train_on_batch 来训练深度学习模型。 我无法将所有行都放入内存中,因此我想一次获得 10K 左右的数据以批量处理成 64 或 128 块(取决于模型大小)。 我目前正在使用 rdd.sample() 但我认为这不能保证我会获得所有行。 是否有更好的方法来分区数据以使其更易于管理,以便我可以编写生成器函数来获取批次? 我的代码如下:
data_df = spark.read.parquet(PARQUET_FILE)
print(f'RDD Count: {data_df.count()}') # 6B+
data_sample = data_df.sample(True, 0.0000015).take(6400)
sample_df = data_sample.toPandas()
def get_batch():
for row in sample_df.itertuples():
# TODO: put together a batch size of BATCH_SIZE
yield row
for i in range(10):
print(next(get_batch()))
我不相信 spark 让您对数据进行偏移或分页。
但是您可以添加一个索引,然后对其进行分页,首先:
from pyspark.sql.functions import lit
data_df = spark.read.parquet(PARQUET_FILE)
count = data_df.count()
chunk_size = 10000
# Just adding a column for the ids
df_new_schema = data_df.withColumn('pres_id', lit(1))
# Adding the ids to the rdd
rdd_with_index = data_df.rdd.zipWithIndex().map(lambda (row,rowId): (list(row) + [rowId+1]))
# Creating a dataframe with index
df_with_index = spark.createDataFrame(chunk_rdd,schema=df_new_schema.schema)
# Iterating into the chunks
for chunk_size in range(0,count+1 ,chunk_size):
initial_page = page_num*chunk_size
final_page = initial_page + chunk_size
where_query = ('pres_id > {0} and pres_id <= {1}').format(initial_page,final_page)
chunk_df = df_with_index.where(where_query).toPandas()
train_on_batch(chunk_df) # <== Your function here
这不是最佳选择,由于使用了熊猫数据框,它会严重利用 spark,但可以解决您的问题。
如果这会影响您的功能,请不要忘记删除 id。
尝试这个:
from pyspark.sql import functions as F
sample_dict = {}
# Read the parquet file
df = spark.read.parquet("parquet file")
# add the partition_number as a column
df = df.withColumn('partition_num', F.spark_partition_id())
df.persist()
total_partition = [int(row.partition_num) for row in
df.select('partition_num').distinct().collect()]
for each_df in total_partition:
sample_dict[each_df] = df.where(df.partition_num == each_df)
我了解到您正计划训练深度学习模型。 看看正是为此用例创建的 Petastorm 开源库。
https://docs.databricks.com/applications/machine-learning/load-data/petastorm.html
Petastorm是一个开源数据访问库。 该库支持直接从 Apache Parquet 格式的数据集和已作为 Apache Spark DataFrames 加载的数据集对深度学习模型进行单节点或分布式训练和评估。 Petastorm 支持流行的基于 Python 的机器学习 (ML) 框架,例如 Tensorflow、PyTorch 和 PySpark。 有关 Petastorm 的更多信息,请参阅 Petastorm GitHub 页面和Petastorm API 文档。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.