按行索引拆分 Spark 数据帧

Posted

技术标签:

【中文标题】按行索引拆分 Spark 数据帧【英文标题】:Split Spark dataframe by row index 【发布时间】:2018-02-20 23:53:46 【问题描述】:

我想按行顺序拆分数据框。如果有 100 行,则希望拆分为 4 个相等的数据帧应分别具有索引 0-2425-4950-7475-99

唯一可用的预定义函数是randomSplit。但是randomSplit 在拆分之前会随机化数据。 我想到的另一种方法是使用count 减少操作查找数据计数,然后使用take 继续提取数据,但这非常昂贵。 有没有其他方法可以在保持相同顺序的同时实现上述目标?

【问题讨论】:

类似:***.com/questions/48884960/… 【参考方案1】:

您可以使用monotonically_increasing_id 获取行号(如果您还没有),然后在行号窗口中使用ntile 来拆分为您想要的任意多个分区:

from pyspark.sql.window import Window
from pyspark.sql.functions import monotonically_increasing_id, ntile

values = [(str(i),) for i in range(100)]
df = spark.createDataFrame(values, ('value',))

def split_by_row_index(df, num_partitions=4):
    # Let's assume you don't have a row_id column that has the row order
    t = df.withColumn('_row_id', monotonically_increasing_id())
    # Using ntile() because monotonically_increasing_id is discontinuous across partitions
    t = t.withColumn('_partition', ntile(num_partitions).over(Window.orderBy(t._row_id))) 
    return [t.filter(t._partition == i+1).drop('_row_id', '_partition') for i in range(partitions)]

[i.collect() for i in split_by_row_index(df)]

【讨论】:

我们如何为更大的df实现相同的?

以上是关于按行索引拆分 Spark 数据帧的主要内容,如果未能解决你的问题,请参考以下文章

使用行分隔符拆分 Spark 数据帧

使用 spark python 拆分数据帧

Spark中的拆分,操作和联合数据框

Pyspark 数据帧拆分并将分隔列值填充到 N 索引数组中

Spark 从超级数据帧优化方法生成子数据帧

如何根据原始数据帧中的总行数将数据帧拆分为两个数据帧